From 0bf4a40077b67627059d107b363f6f700a7bbf29 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 20 Aug 2026 11:16:29 +0100 Subject: [PATCH 01/21] Cost linear map preservation against expansion Expanding a pullback exposes scalar factorisation across its entries; preserving it exposes a physical basis that both argument axes share. Neither dominates: preserving wins on the Piola mapped families, where the geometry otherwise crosses the element tensor contraction twice, and expanding wins on Lagrange at moderate degree, where the entries carry enough structure to fold. Factorisation is therefore parameterised on that choice and run twice, and the cheaper plan by estimate_cost is kept. The second run is skipped when no sum spans exactly one argument axis, which is every tensor product cell here: sum factorisation has already split the basis into one dimensional factors and contracted the Jacobian into the per point geometry, so no mapped tabulation exists to share and the search would cost compile time for an identical plan. Preserved maps must survive finalisation to be shared at all, so spectral mode keeps its ComponentTensors. Measured on inner(u, v)*dx + inner(d(u), d(v))*dx, d the family's derivative, against the same tree without this change: RT tetrahedra degree 5 12,965,694 -> 8,711,463 flops, 33% fewer RT tetrahedra degree 3 374,589 -> 255,933 flops, 32% fewer RT triangles degree 5 253,051 -> 192,564 flops, 24% fewer CG tetrahedra degree 1 432 -> 390 flops, 10% fewer Q, NCE hexahedra unchanged, no map to preserve Scalar temporaries fall 31% and AST lines 11% on RT tetrahedra, and no case regresses on flops. Compile time rises by up to 17% where a second plan is built, and is unchanged elsewhere. Fewer operations do not yet make a faster kernel. The RT bilinear kernel runs 11% slower on triangles and 4.5% slower on tetrahedra at degree 3, reproducibly: each shared map becomes its own ComponentTensor, and scheduling gives each one its own loop, so one fused loop over the basis becomes three over the same extent. Stacking the maps that share an extent into one tensor is what this needs next. Co-Authored-By: Claude Opus 5 --- tests/tsfc/test_impero_loopy_flop_counts.py | 21 +++- tests/tsfc/test_sum_factorisation.py | 50 ++++++++- tsfc/spectral.py | 114 +++++++++++++++----- 3 files changed, 153 insertions(+), 32 deletions(-) diff --git a/tests/tsfc/test_impero_loopy_flop_counts.py b/tests/tsfc/test_impero_loopy_flop_counts.py index 240067bd6e..4bbdef0dc6 100644 --- a/tests/tsfc/test_impero_loopy_flop_counts.py +++ b/tests/tsfc/test_impero_loopy_flop_counts.py @@ -6,9 +6,9 @@ import loopy from tsfc import compile_form from ufl import (FunctionSpace, Mesh, TestFunction, - TrialFunction, dx, grad, inner, + TrialFunction, div, dx, grad, inner, interval, triangle, quadrilateral, - TensorProductCell) + tetrahedron, TensorProductCell) from finat.ufl import FiniteElement, VectorElement from tsfc.parameters import target @@ -64,3 +64,20 @@ def test_flop_count(cell, parameters): loopy_flops = numpy.asarray(loopy_flops) assert all(new_flops == loopy_flops) + + +@pytest.mark.parametrize("cell", [triangle, tetrahedron], + ids=lambda cell: cell.cellname) +def test_flop_count_mapped_tabulation(cell): + # Preserving a Piola map materialises it as a ComponentTensor. + # Scheduling emits that as an assignment, not a loop nest, so counting + # it needs its own extents. + mesh = Mesh(VectorElement("P", cell, 1)) + for k in range(1, 4): + V = FunctionSpace(mesh, FiniteElement("RT", cell, k)) + u = TrialFunction(V) + v = TestFunction(V) + a = inner(u, v)*dx + inner(div(u), div(v))*dx + kernel, = compile_form(a, prefix="form", + parameters={"mode": "spectral"}) + assert kernel.flop_count == count_loopy_flops(kernel) diff --git a/tests/tsfc/test_sum_factorisation.py b/tests/tsfc/test_sum_factorisation.py index 44c1d46dda..f40ffb0fb1 100644 --- a/tests/tsfc/test_sum_factorisation.py +++ b/tests/tsfc/test_sum_factorisation.py @@ -3,10 +3,12 @@ from ufl import (Mesh, FunctionSpace, TestFunction, TrialFunction, TensorProductCell, dx, action, interval, triangle, - quadrilateral, curl, dot, div, grad) + quadrilateral, tetrahedron, curl, dot, div, + grad, inner) from finat.ufl import (FiniteElement, VectorElement, EnrichedElement, TensorProductElement, HCurlElement, HDivElement) +import tsfc.spectral from tsfc import compile_form @@ -190,6 +192,52 @@ def test_vector_laplace_action(cell, order): assert (rates < order).all() +@pytest.fixture +def expanded(monkeypatch): + """Force the expanded representation, for comparison.""" + def force(monkeypatch=monkeypatch): + collect_monomials = tsfc.spectral.collect_monomials + monkeypatch.setattr( + tsfc.spectral, "collect_monomials", + lambda expressions, classifier, _: collect_monomials( + expressions, classifier)) + return force + + +def piola_helmholtz(cell, degree): + m = Mesh(VectorElement('CG', cell, 1)) + V = FunctionSpace(m, FiniteElement('RT', cell, degree)) + u = TrialFunction(V) + v = TestFunction(V) + return (inner(u, v) + inner(div(u), div(v)))*dx + + +@pytest.mark.parametrize('cell', [triangle, tetrahedron], + ids=lambda cell: cell.cellname) +@pytest.mark.parametrize('degree', [1, 2, 3]) +def test_piola_map_is_preserved(cell, degree, expanded): + # Test and trial apply the same Piola map, so preserving it evaluates + # the physical basis once instead of pushing the geometry through both + # argument axes. + form = piola_helmholtz(cell, degree) + selected = count_flops(form) + expanded() + assert selected < count_flops(form) + + +@pytest.mark.parametrize('cell', [triangle, tetrahedron], + ids=lambda cell: cell.cellname) +@pytest.mark.parametrize('degree', [1, 3]) +def test_preserving_a_map_is_never_worse(cell, degree, expanded): + # Expanding a map exposes scalar factorisation of its entries, which at + # some degrees beats sharing it. Selection costs both, so neither + # representation may regress the other. + form = helmholtz(cell, degree) + selected = count_flops(form) + expanded() + assert selected <= count_flops(form) + + if __name__ == "__main__": import os import sys diff --git a/tsfc/spectral.py b/tsfc/spectral.py index a521fdb2fd..0f76034da0 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -4,7 +4,8 @@ from gem.gem import Delta, Indexed, Sum, index_sum, one from gem.node import Memoizer, MemoizerArg -from gem.optimise import filtered_replace_indices +from gem.optimise import (estimate_cost, filtered_replace_indices, + has_linear_maps) from gem.optimise import delta_elimination as _delta_elimination from gem.optimise import replace_division, unroll_indexsum from gem.refactorise import ATOMIC, COMPOUND, OTHER, MonomialSum, collect_monomials @@ -52,34 +53,50 @@ def _delta_inside(node, self): for child in node.children) -def flatten(var_reps, index_cache): - quadrature_indices = OrderedDict() - - pairs = [] # assignment pairs - for variable, reps in var_reps: - # Extract argument indices - argument_indices, = set(r.argument_indices for r in reps) - assert set(variable.free_indices) == set(argument_indices) +def _group_key(pair): + variable, expression = pair + return frozenset(variable.free_indices) - # Extract and verify expressions - expressions = [r.expression for r in reps] - assert all(set(e.free_indices) <= set(argument_indices) - for e in expressions) - # Save assignment pair - pairs.append((variable, Sum(*expressions))) +def _preservable(pairs): + """Is there a linear map whose preservation could change a plan? - # Collect quadrature_indices - for r in reps: - quadrature_indices.update(zip_longest(r.quadrature_multiindex, ())) + Parameters + ---------- + pairs : tuple of tuple + Output variables and the integrands assigned to them. - # Split Concatenate nodes - pairs = unconcatenate(pairs, cache=index_cache) + Returns + ------- + bool + Whether any assignment contains a sum over one argument axis. - def group_key(pair): - variable, expression = pair - return frozenset(variable.free_indices) + """ + return any( + has_linear_maps([expression], set(free_indices)) + for free_indices, pair_group in groupby(pairs, _group_key) + for _, expression in pair_group) + + +def _factorise(pairs, quadrature_indices, preserve_maps): + """Factorise the arguments of each assignment and place its reductions. + + Parameters + ---------- + pairs : tuple of tuple + Output variables and the integrands assigned to them. + quadrature_indices : tuple of Index + Every quadrature index of the integral, in source order. + preserve_maps : bool + Keep a sum over one argument axis whole, as a finite element linear + map, rather than distributing it into scalar monomials. + + Returns + ------- + tuple of tuple + Output variables and their factorised GEM expressions. + """ # Common memoizer to remove ComponentTensors index_replacer = MemoizerArg(filtered_replace_indices) # Common memoizer to test for Deltas inside expressions @@ -89,11 +106,14 @@ def group_key(pair): # Assignments are variable -> MonomialSum map delta_simplified = defaultdict(MonomialSum) # Group assignment pairs by argument indices - for free_indices, pair_group in groupby(pairs, group_key): + for free_indices, pair_group in groupby(pairs, _group_key): variables, expressions = zip(*pair_group) + argument_indices = set(free_indices) + classifier = partial(classify, argument_indices, delta_inside=delta_inside) # Argument factorise expressions - classifier = partial(classify, set(free_indices), delta_inside=delta_inside) - monomial_sums = collect_monomials(expressions, classifier) + monomial_sums = collect_monomials( + expressions, classifier, + argument_indices if preserve_maps else ()) # For each monomial, apply delta cancellation and insert # result into delta_simplified. for variable, monomial_sum in zip(variables, monomial_sums): @@ -103,6 +123,7 @@ def group_key(pair): delta_simplified[var].add(s, a, r) # Final factorisation + plan = [] for variable in narrow_variables: monomial_sum = delta_simplified[variable] # Collect sum indices applicable to the current MonomialSum @@ -110,11 +131,46 @@ def group_key(pair): # 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 = sum_factorise(variable, sum_indices, monomial_sum) - yield (variable, expression) + plan.append((variable, sum_factorise(variable, sum_indices, monomial_sum))) + return tuple(plan) + + +def flatten(var_reps, index_cache): + quadrature_indices = OrderedDict() + + pairs = [] # assignment pairs + for variable, reps in var_reps: + # Extract argument indices + argument_indices, = set(r.argument_indices for r in reps) + assert set(variable.free_indices) == set(argument_indices) + + # Extract and verify expressions + expressions = [r.expression for r in reps] + assert all(set(e.free_indices) <= set(argument_indices) + for e in expressions) + + # Save assignment pair + pairs.append((variable, Sum(*expressions))) + + # Collect quadrature_indices + for r in reps: + quadrature_indices.update(zip_longest(r.quadrature_multiindex, ())) + + # Split Concatenate nodes + pairs = unconcatenate(pairs, cache=index_cache) + + # Expanding a linear map exposes scalar factorisation across its + # entries, preserving one exposes a tabulation that several argument + # axes share, and neither dominates. Cost both and keep the cheaper, + # skipping the second factorisation when there is no map to preserve. + plans = [_factorise(pairs, quadrature_indices, False)] + if _preservable(pairs): + plans.append(_factorise(pairs, quadrature_indices, True)) + return min(plans, key=lambda plan: estimate_cost( + expression for _, expression in plan)) -finalise_options = dict(replace_delta=False) +finalise_options = dict(replace_delta=False, remove_componenttensors=False) def classify(argument_indices, expression, delta_inside): From bfa2a9b54023e57835ebb0f77fa026f176b968cb Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 20 Aug 2026 13:26:08 +0100 Subject: [PATCH 02/21] Tabulate adjacent shared maps in one loop A preserved linear map is materialised as a ComponentTensor, and the loopy backend minted a fresh iname for each one, so several maps over one extent became several loops where expansion emits a single fused nest. That fission cost 7-11% of the bilinear kernel. Reuse the iname between tabulations that the schedule places side by side. Only adjacent ones: impero interleaves statements that depend on a tabulation, and one iname can not sit both inside and outside such a statement, which loopy reports as a scheduling cycle. Co-Authored-By: Claude Opus 5 --- tests/tsfc/test_sum_factorisation.py | 19 +++++++++++++++++++ tsfc/loopy.py | 27 ++++++++++++++++++++++++--- 2 files changed, 43 insertions(+), 3 deletions(-) diff --git a/tests/tsfc/test_sum_factorisation.py b/tests/tsfc/test_sum_factorisation.py index f40ffb0fb1..0a7e5ad4d7 100644 --- a/tests/tsfc/test_sum_factorisation.py +++ b/tests/tsfc/test_sum_factorisation.py @@ -81,6 +81,11 @@ def count_storage(form): and temporary.initializer is not None)) +def count_loops(form): + kernel, = compile_form(form, parameters=dict(mode='spectral')) + return len(kernel.ast.default_entrypoint.all_inames()) + + @pytest.mark.parametrize(('cell', 'order'), [(quadrilateral, 5), (TensorProductCell(interval, interval), 5), @@ -238,6 +243,20 @@ def test_preserving_a_map_is_never_worse(cell, degree, expanded): assert selected <= count_flops(form) +@pytest.mark.parametrize('cell', [triangle, tetrahedron], + ids=lambda cell: cell.cellname) +@pytest.mark.parametrize('degree', [1, 3]) +def test_shared_map_is_tabulated_in_one_loop(cell, degree, expanded): + # A map that both argument axes share is tabulated once, so it must be + # tabulated in one loop. An index per axis fissions the loop nest that + # the expanded representation keeps whole, which costs more than the + # flops it saves. + form = piola_helmholtz(cell, degree) + selected = count_loops(form) + expanded() + assert selected <= count_loops(form) + + if __name__ == "__main__": import os import sys diff --git a/tsfc/loopy.py b/tsfc/loopy.py index d4a31a36cb..0c2d1f12bc 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.tabulated = (None, ()) # axes and inames of the preceding tabulation self.gem_to_pymbolic = {} # gem node -> pymbolic variable self.name_gen = UniqueNameGenerator() self.target = target @@ -304,10 +305,27 @@ def statement(tree, ctx): raise AssertionError("cannot generate loopy from %s" % type(tree)) +def tabulated_axes(tree): + """The axes a statement binds, if it tabulates a tensor over its own.""" + if isinstance(tree, imp.Evaluate) \ + and isinstance(tree.expression, gem.ComponentTensor): + return tree.expression.multiindex + return None + + @statement.register(imp.Block) def statement_block(tree, ctx): - from itertools import chain - return list(chain(*(statement(child, ctx) for child in tree.children))) + # Tabulations of the same axes share a loop while they stay adjacent. + # Anything between them is a statement the schedule placed outside that + # loop, so the loop has to close before it and reopen after. + instructions = [] + ctx.tabulated = (None, ()) + for child in tree.children: + instructions.extend(statement(child, ctx)) + if tabulated_axes(child) is None: + ctx.tabulated = (None, ()) + ctx.tabulated = (None, ()) + return instructions @statement.register(imp.For) @@ -360,7 +378,10 @@ def statement_evaluate(leaf, ctx): elif isinstance(expr, gem.Constant): return [] elif isinstance(expr, gem.ComponentTensor): - idx = ctx.gem_to_pym_multiindex(expr.multiindex) + axes, idx = ctx.tabulated + if axes != expr.multiindex: + idx = ctx.gem_to_pym_multiindex(expr.multiindex) + ctx.tabulated = (expr.multiindex, idx) 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: From 2f711491adb81f1c99995e00fc94927477405798 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 20 Aug 2026 16:10:46 +0100 Subject: [PATCH 03/21] DROP BEFORE MERGE: build FIAT from the PR stack The TSFC changes here need the GEM changes in the FIAT stack firedrakeproject/fiat#282 -> #284 -> #281 -> #286, whose head carries all four. Install it over the one pyproject.toml resolves from main, so that CI exercises both halves together. Revert this commit once the FIAT stack lands. Co-Authored-By: Claude Opus 5 --- .github/actions/install/action.yml | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/.github/actions/install/action.yml b/.github/actions/install/action.yml index 54c1411d7f..5a55c5aed0 100644 --- a/.github/actions/install/action.yml +++ b/.github/actions/install/action.yml @@ -163,6 +163,17 @@ runs: firedrake-clean pip list + - name: 'DROP BEFORE MERGE: build FIAT from the PR stack' + shell: bash + run: | + . venv/bin/activate + : # Head of the stack firedrakeproject/fiat#282 -> #284 -> #281 -> #286 + pip install --no-deps --force-reinstall \ + git+https://github.com/firedrakeproject/fiat.git@pbrubeck/coffee-scalar-factor + : # The code generators changed, so discard kernels built against main + firedrake-clean + pip list + - name: Run firedrake-check shell: bash run: | From 29a7ec023ee3e4e1da5a8ed7eaba0bcc8c33357f Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 20 Aug 2026 18:52:53 +0100 Subject: [PATCH 04/21] Apply suggestion from @pbrubeck --- .github/actions/install/action.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/actions/install/action.yml b/.github/actions/install/action.yml index 5a55c5aed0..73264ebca7 100644 --- a/.github/actions/install/action.yml +++ b/.github/actions/install/action.yml @@ -168,7 +168,7 @@ runs: run: | . venv/bin/activate : # Head of the stack firedrakeproject/fiat#282 -> #284 -> #281 -> #286 - pip install --no-deps --force-reinstall \ + pip install --no-deps --force-reinstall --ignore-installed \ git+https://github.com/firedrakeproject/fiat.git@pbrubeck/coffee-scalar-factor : # The code generators changed, so discard kernels built against main firedrake-clean From 5bc40c93e193efabfc84f59ba7874574c84caf1e Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 21 Aug 2026 15:15:00 +0100 Subject: [PATCH 05/21] Lower indirect reduction views without copies --- tsfc/loopy.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/tsfc/loopy.py b/tsfc/loopy.py index 0c2d1f12bc..43aee0610f 100644 --- a/tsfc/loopy.py +++ b/tsfc/loopy.py @@ -568,7 +568,18 @@ def _expression_variable(expr, ctx): @_expression.register(gem.Indexed) def _expression_indexed(expr, ctx): rank = ctx.fetch_multiindex(expr.multiindex) - var = expression(expr.children[0], ctx) + aggregate, = expr.children + if (isinstance(aggregate, gem.ComponentTensor) + and aggregate not in ctx.gem_to_pymbolic): + body, = aggregate.children + if body in ctx.gem_to_pymbolic: + replacements = dict(zip(aggregate.multiindex, expr.multiindex)) + multiindex = tuple(replacements.get(index, index) + for index in ctx.indices[body]) + rank = ctx.fetch_multiindex(multiindex) + return p.Subscript(ctx._gem_to_pym_var(body), rank) + + var = expression(aggregate, ctx) if isinstance(var, p.Subscript): rank = var.index + rank var = var.aggregate From e59065dc911852b0684458bf74b8fde061b24863 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 22 Aug 2026 16:55:44 +0100 Subject: [PATCH 06/21] Place indirect reductions once per assignment GEM no longer factors reductions through indirect gathers inside optimise_monomial_sum, which the recursive sum_factorise calls at every level. Apply the traversal once to each finished assignment instead, so plan costing still sees its effect. Co-Authored-By: Claude Opus 5 --- tsfc/coffee_mode.py | 6 ++++-- tsfc/spectral.py | 10 ++++++---- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/tsfc/coffee_mode.py b/tsfc/coffee_mode.py index 632b915b41..cd5d383fa7 100644 --- a/tsfc/coffee_mode.py +++ b/tsfc/coffee_mode.py @@ -2,7 +2,8 @@ from gem.node import traversal, Memoizer from gem.gem import Failure, Sum, index_sum -from gem.optimise import replace_division, unroll_indexsum +from gem.optimise import (factorise_indirect_reductions, replace_division, + unroll_indexsum) from gem.refactorise import collect_monomials from gem.unconcatenate import unconcatenate from gem.coffee import optimise_monomial_sum @@ -78,4 +79,5 @@ def optimise_expressions(expressions, argument_indices): classifier = partial(spectral.classify, set(argument_indices), delta_inside=Memoizer(spectral._delta_inside)) monomial_sums = collect_monomials(expressions, classifier) - return [optimise_monomial_sum(ms, argument_indices) for ms in monomial_sums] + return [factorise_indirect_reductions( + optimise_monomial_sum(ms, argument_indices)) for ms in monomial_sums] diff --git a/tsfc/spectral.py b/tsfc/spectral.py index 0f76034da0..e261a1aa1e 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -4,8 +4,8 @@ from gem.gem import Delta, Indexed, Sum, index_sum, one from gem.node import Memoizer, MemoizerArg -from gem.optimise import (estimate_cost, filtered_replace_indices, - has_linear_maps) +from gem.optimise import (estimate_cost, factorise_indirect_reductions, + filtered_replace_indices, has_linear_maps) from gem.optimise import delta_elimination as _delta_elimination from gem.optimise import replace_division, unroll_indexsum from gem.refactorise import ATOMIC, COMPOUND, OTHER, MonomialSum, collect_monomials @@ -130,8 +130,10 @@ def _factorise(pairs, quadrature_indices, preserve_maps): 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 - plan.append((variable, sum_factorise(variable, sum_indices, monomial_sum))) + # Apply sum factorisation combined with COFFEE technology, then + # place each reduction against the whole factorised assignment. + expression = sum_factorise(variable, sum_indices, monomial_sum) + plan.append((variable, factorise_indirect_reductions(expression))) return tuple(plan) From aa99e04a43f16a71ab106da308aaa61bb603a67e Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sun, 23 Aug 2026 11:39:57 +0100 Subject: [PATCH 07/21] Cancel the Deltas that select basis transformation columns A basis transformation is now a contraction against a Delta, and delta_elimination only inspects top-level product factors, so the Delta inside a preserved linear map never reaches it. Cancel those before monomial collection, which recovers the gather the transformation used to build directly. Co-Authored-By: Claude Opus 5 --- tsfc/spectral.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tsfc/spectral.py b/tsfc/spectral.py index e261a1aa1e..5f7478dc39 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -4,7 +4,8 @@ from gem.gem import Delta, Indexed, Sum, index_sum, one from gem.node import Memoizer, MemoizerArg -from gem.optimise import (estimate_cost, factorise_indirect_reductions, +from gem.optimise import (eliminate_deltas, estimate_cost, + factorise_indirect_reductions, filtered_replace_indices, has_linear_maps) from gem.optimise import delta_elimination as _delta_elimination from gem.optimise import replace_division, unroll_indexsum @@ -35,6 +36,10 @@ def Integrals(expressions, quadrature_multiindex, argument_multiindices, paramet # Rewrite: a / b => a * (1 / b) expressions = replace_division(expressions) + # Cancel the Deltas that select a basis transformation's columns, so that + # monomial collection sees the resulting gather rather than the Delta. + expressions = [eliminate_deltas(e) for e in expressions] + # Unroll max_extent = parameters["unroll_indexsum"] if max_extent: From 5e7080bad70f2a440ca8adea0b857561c64c106a Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 24 Aug 2026 17:55:21 +0100 Subject: [PATCH 08/21] Follow the gem.optimise rename to tabulate_indirect_contractions Co-Authored-By: Claude Opus 5 --- tsfc/coffee_mode.py | 4 ++-- tsfc/spectral.py | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/tsfc/coffee_mode.py b/tsfc/coffee_mode.py index cd5d383fa7..845b52ed1c 100644 --- a/tsfc/coffee_mode.py +++ b/tsfc/coffee_mode.py @@ -2,7 +2,7 @@ from gem.node import traversal, Memoizer from gem.gem import Failure, Sum, index_sum -from gem.optimise import (factorise_indirect_reductions, replace_division, +from gem.optimise import (tabulate_indirect_contractions, replace_division, unroll_indexsum) from gem.refactorise import collect_monomials from gem.unconcatenate import unconcatenate @@ -79,5 +79,5 @@ def optimise_expressions(expressions, argument_indices): classifier = partial(spectral.classify, set(argument_indices), delta_inside=Memoizer(spectral._delta_inside)) monomial_sums = collect_monomials(expressions, classifier) - return [factorise_indirect_reductions( + return [tabulate_indirect_contractions( optimise_monomial_sum(ms, argument_indices)) for ms in monomial_sums] diff --git a/tsfc/spectral.py b/tsfc/spectral.py index 5f7478dc39..bb9665eefa 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -5,7 +5,7 @@ from gem.gem import Delta, Indexed, Sum, index_sum, one from gem.node import Memoizer, MemoizerArg from gem.optimise import (eliminate_deltas, estimate_cost, - factorise_indirect_reductions, + tabulate_indirect_contractions, filtered_replace_indices, has_linear_maps) from gem.optimise import delta_elimination as _delta_elimination from gem.optimise import replace_division, unroll_indexsum @@ -136,9 +136,9 @@ def _factorise(pairs, quadrature_indices, preserve_maps): # 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, then - # place each reduction against the whole factorised assignment. + # tabulate indirect contractions over the whole factorised assignment. expression = sum_factorise(variable, sum_indices, monomial_sum) - plan.append((variable, factorise_indirect_reductions(expression))) + plan.append((variable, tabulate_indirect_contractions(expression))) return tuple(plan) From 9f00670d407a9cdd40193b7f2f92b4f3739fafd1 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 17 Aug 2026 13:06:38 +0100 Subject: [PATCH 09/21] Lower jagged contraction domains in Loopy A jagged index is bounded by its parents, so its ISL domain must be built against them rather than as an independent axis. Loopy generation now carries the parent inames alongside the extents and constrains each dependent index inside the loops that bound it, which is what lets a sparse basis map and a simplex lattice reach the generated kernel without a data-dependent loop bound. Co-Authored-By: Claude Opus 5 --- tests/tsfc/test_pickle_gem.py | 14 +++++++++ tsfc/loopy.py | 59 +++++++++++++++++++++++++++-------- 2 files changed, 60 insertions(+), 13 deletions(-) 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 43aee0610f..1aaface191 100644 --- a/tsfc/loopy.py +++ b/tsfc/loopy.py @@ -124,6 +124,7 @@ def __init__(self, target=None): self.active_indices = {} # gem index -> pymbolic variable self.index_extent = OrderedDict() # pymbolic variable for indices -> extent self.tabulated = (None, ()) # axes and inames of the preceding tabulation + self.index_parents = {} # iname -> parent inames bounding a jagged index self.gem_to_pymbolic = {} # gem node -> pymbolic variable self.name_gen = UniqueNameGenerator() self.target = target @@ -258,7 +259,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( @@ -277,16 +278,30 @@ def generate(impero_c, args, scalar_type, kernel_name="loopy_kernel", index_name return knl, event_name -def create_domains(indices): - """ Create ISL domains from indices +def create_domains(indices, index_parents=None): + """Create ISL domains for independent and dependent indices. - :arg indices: iterable of (index_name, extent) pairs - :returns: A list of ISL sets representing the iteration domain of the indices.""" + Parameters + ---------- + indices : iterable of tuple + Index names and their static extents. + index_parents : mapping, optional + Parent inames for simplex-lattice bounds. + Returns + ------- + list of isl.Set + Iteration domains for Loopy. + """ domains = [] for idx, extent in indices: - inames = isl.make_zero_and_vars([idx]) - domains.append(((inames[0].le_set(inames[idx])) & (inames[idx].lt_set(inames[0] + extent)))) + parents = index_parents.get(idx, ()) if index_parents else () + inames = isl.make_zero_and_vars([idx], parents) + bound = inames[0] + extent + for parent in parents: + bound = bound - inames[parent] + domains.append(inames[0].le_set(inames[idx]) + & inames[idx].lt_set(bound)) if not domains: domains = [isl.BasicSet("[] -> {[]}")] @@ -334,6 +349,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) @@ -378,13 +400,24 @@ def statement_evaluate(leaf, ctx): elif isinstance(expr, gem.Constant): return [] elif isinstance(expr, gem.ComponentTensor): - axes, idx = ctx.tabulated - if axes != expr.multiindex: - idx = ctx.gem_to_pym_multiindex(expr.multiindex) - ctx.tabulated = (expr.multiindex, idx) + implicit_axes = tuple( + index for index in expr.multiindex + if index not in ctx.active_indices) + axes, implicit_values = ctx.tabulated + if axes != implicit_axes: + implicit_values = ctx.gem_to_pym_multiindex(implicit_axes) + ctx.tabulated = (implicit_axes, implicit_values) + implicit_indices = dict(zip(implicit_axes, implicit_values)) + value_indices = [] + for index in expr.multiindex: + if index in ctx.active_indices: + value_indices.append(ctx.active_indices[index]) + else: + value_indices.append(implicit_indices[index]) + value_indices = tuple(value_indices) var, sub_idx = ctx.pymbolic_variable_and_destruct(expr) - lhs = p.Subscript(var, sub_idx + idx) - with active_indices(dict(zip(expr.multiindex, idx)), ctx) as ctx_active: + lhs = p.Subscript(var, sub_idx + value_indices) + with active_indices(implicit_indices, ctx) as ctx_active: return [lp.Assignment(lhs, expression(expr.children[0], ctx_active), within_inames=ctx_active.active_inames())] elif isinstance(expr, gem.Inverse): idx = ctx.pymbolic_multiindex(expr.shape) From 186490384459df1bb2865571f6eaffd92d6694d6 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 13 Aug 2026 22:50:58 +0100 Subject: [PATCH 10/21] Compact products of simplex lattice temporaries --- tsfc/loopy.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/tsfc/loopy.py b/tsfc/loopy.py index 1aaface191..68f0fc7f28 100644 --- a/tsfc/loopy.py +++ b/tsfc/loopy.py @@ -126,6 +126,7 @@ def __init__(self, target=None): self.tabulated = (None, ()) # axes and inames of the preceding tabulation self.index_parents = {} # iname -> parent inames bounding a jagged index self.gem_to_pymbolic = {} # gem node -> pymbolic variable + self.compact_indices = {} # temporary -> compact index layout self.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 From 422ffbce0a0ec837e2f452709b8421a8a4da7a14 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sun, 2 Aug 2026 15:15:08 +0100 Subject: [PATCH 11/21] WIP: layer simplex sum-factorisation tests --- pyproject.toml | 3 +- tests/firedrake/regression/test_quadrature.py | 39 +++++++ tests/tsfc/test_codegen.py | 53 ++++++++- tests/tsfc/test_sum_factorisation.py | 109 ++++++++++++++++++ tsfc/kernel_interface/common.py | 10 +- 5 files changed, 211 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index ecbeaaa735..e5e03a13dc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,8 @@ dependencies = [ # each Firedrake release to a specific UFL minor version (e.g. 2025.3.x) "fenics-ufl @ git+https://github.com/FEniCS/ufl.git@main", # TODO RELEASE - "firedrake-fiat @ git+https://github.com/firedrakeproject/fiat.git@main", + # DROP BEFORE MERGE: pinned to the paired FIAT branch for CI; revert to @main + "firedrake-fiat @ git+https://github.com/firedrakeproject/fiat.git@pbrubeck/simplex-sum-factor", "h5py>3.12.1", "firedrake-rtree>=2026.2.0", "immutabledict", diff --git a/tests/firedrake/regression/test_quadrature.py b/tests/firedrake/regression/test_quadrature.py index 225a4b244d..bb7189ddc4 100644 --- a/tests/firedrake/regression/test_quadrature.py +++ b/tests/firedrake/regression/test_quadrature.py @@ -52,3 +52,42 @@ def test_quadrature_element(mesh, family, mat_type, diagonal): a = inner(u, v) * dx assemble(a, mat_type=mat_type, diagonal=diagonal) + + +@pytest.mark.parametrize("family", ["DG", "CG", "Bernstein"]) +@pytest.mark.parametrize("cell", ["triangle", "tetrahedron"]) +@pytest.mark.parametrize("degree", [1, 3]) +def test_collapsed_quadrature_sum_factorisation(cell, degree, family): + """Check sum-factorized residuals and matrices against dense tabulation.""" + mesh = {"triangle": UnitSquareMesh(2, 2), + "tetrahedron": UnitCubeMesh(1, 1, 1)}[cell] + variant = None if family == "Bernstein" else "integral" + V = FunctionSpace(mesh, family, degree, variant=variant) + u = TrialFunction(V) + v = TestFunction(V) + rg = RandomGenerator(PCG64(seed=0)) + w = rg.uniform(V, 0, 1) + + # translate_coefficient path (forward transform): residual with a + # derivative, mixing both the coefficient and argument sum-factorized + # tabulations. + L = inner(grad(w), grad(v)) * dx(scheme="canonical") + L_collapsed = inner(grad(w), grad(v)) * dx(scheme="collapsed") + b = assemble(L) + b_collapsed = assemble(L_collapsed) + assert np.allclose(b.dat.data, b_collapsed.dat.data, rtol=1e-10, atol=1e-10) + + # translate_argument path (backward transform): mass matrix. + a = inner(u, v) * dx(scheme="canonical") + a_collapsed = inner(u, v) * dx(scheme="collapsed") + M = assemble(a).M.values + M_collapsed = assemble(a_collapsed).M.values + assert np.allclose(M, M_collapsed, rtol=1e-10, atol=1e-10) + + # Bilinear derivatives exercise two independently transformed argument + # lattices. + a = inner(grad(u), grad(v)) * dx(scheme="canonical") + a_collapsed = inner(grad(u), grad(v)) * dx(scheme="collapsed") + K = assemble(a).M.values + K_collapsed = assemble(a_collapsed).M.values + assert np.allclose(K, K_collapsed, rtol=1e-10, atol=1e-10) diff --git a/tests/tsfc/test_codegen.py b/tests/tsfc/test_codegen.py index 8d0bc79655..7dd97e2ba1 100644 --- a/tests/tsfc/test_codegen.py +++ b/tests/tsfc/test_codegen.py @@ -1,6 +1,7 @@ +import numpy import pytest -from gem import impero_utils +from gem import gem, impero_utils from gem.gem import Index, Indexed, IndexSum, Product, Variable @@ -24,6 +25,56 @@ def gencode(expr): assert len(gencode(e1).children) == len(gencode(e2).children) +def test_jagged_index_codegen(monkeypatch): + import islpy as isl + import loopy as lp + import tsfc.loopy + + # Execute the generated code so we check the numbers, not just the loop bounds + monkeypatch.setattr(tsfc.loopy, "target", lp.ExecutableCTarget()) + + n = 4 + extent = n + 1 + npts = 3 + ndof = (n + 1) * (n + 2) // 2 + + rng = numpy.random.default_rng(7) + # Table zero-padded outside the simplex lattice p + q > n, and a + # clamped Morton index table for the coefficient gather + B = rng.random((extent, extent, npts)) + morton = numpy.zeros((extent, extent), dtype=gem.uint_type) + for p_, q_ in numpy.ndindex(morton.shape): + if p_ + q_ > n: + B[p_, q_] = 0.0 + else: + morton[p_, q_] = (p_ + q_) * (p_ + q_ + 1) // 2 + q_ + c = rng.random(ndof) + + i = Index(name="i", extent=npts) + p = Index(name="p", extent=extent) + q = gem.JaggedIndex(name="q", extent=extent, parents=(p,)) + + dof = gem.VariableIndex(Indexed(gem.Literal(morton, dtype=gem.uint_type), (p, q))) + integrand = Product(Indexed(Variable("c", (ndof,)), (dof,)), + Indexed(gem.Literal(B), (p, q, i))) + expr = IndexSum(integrand, (p, q)) + + u = Variable("u", (npts,)) + impero_c = impero_utils.compile_gem([(Indexed(u, (i,)), expr)], (i, p, q)) + args = [lp.GlobalArg("u", dtype=numpy.float64, shape=(npts,)), + lp.GlobalArg("c", dtype=numpy.float64, shape=(ndof,))] + knl, _ = tsfc.loopy.generate(impero_c, args, numpy.float64) + + # The jagged loop must have a domain parametrized by its parent iname + assert any(dom.get_var_names(isl.dim_type.param) + for dom in knl.default_entrypoint.domains) + + u_out = numpy.zeros(npts) + knl(c=c, u=u_out) + u_ref = numpy.tensordot(c[morton], B, axes=((0, 1), (0, 1))) + assert numpy.allclose(u_out, u_ref, rtol=1e-14) + + if __name__ == "__main__": import os import sys diff --git a/tests/tsfc/test_sum_factorisation.py b/tests/tsfc/test_sum_factorisation.py index 0a7e5ad4d7..1bf92871ba 100644 --- a/tests/tsfc/test_sum_factorisation.py +++ b/tests/tsfc/test_sum_factorisation.py @@ -257,6 +257,115 @@ def test_shared_map_is_tabulated_in_one_loop(cell, degree, expanded): assert selected <= count_loops(form) +def simplex_mass(cell, family, degree, scheme='collapsed'): + m = Mesh(VectorElement('CG', cell, 1)) + variant = None if family == "Bernstein" else "integral" + V = FunctionSpace(m, FiniteElement(family, cell, degree, variant=variant)) + u = TrialFunction(V) + v = TestFunction(V) + return inner(u, v) * dx(scheme=scheme) + + +def simplex_laplacian(cell, family, degree, scheme='collapsed'): + m = Mesh(VectorElement('CG', cell, 1)) + variant = None if family == "Bernstein" else "integral" + V = FunctionSpace(m, FiniteElement(family, cell, degree, variant=variant)) + u = TrialFunction(V) + v = TestFunction(V) + return inner(grad(u), grad(v)) * dx(scheme=scheme) + + +@pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"]) +@pytest.mark.parametrize(('cell', 'order'), [(triangle, 3), (tetrahedron, 4)]) +def test_simplex_mass_action(cell, family, order): + degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8)) + flops = [count_flops(action(simplex_mass(cell, family, degree))) + for degree in degrees] + rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees)) + assert (rates < order).all() + + +@pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"]) +@pytest.mark.parametrize(('cell', 'order'), [(triangle, 3), (tetrahedron, 4.4)]) +def test_simplex_laplacian_action(cell, family, order): + degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8)) + flops = [count_flops(action(simplex_laplacian(cell, family, degree))) + for degree in degrees] + rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees)) + assert (rates < order).all() + + +@pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"]) +def test_simplex_laplacian_action_compact_codegen(family): + form = action(simplex_laplacian(triangle, family, 3)) + kernel, = compile_form(form, parameters=dict(mode='spectral')) + temporaries = kernel.ast.default_entrypoint.temporary_variables + assert len(temporaries) < 100 + + +def test_bernstein_laplacian_action_compact_literals(): + degree = 5 + form = action(simplex_laplacian(tetrahedron, "Bernstein", degree)) + kernel, = compile_form(form, parameters=dict(mode='spectral')) + temporaries = kernel.ast.default_entrypoint.temporary_variables + literals = [numpy.asarray(temporary.initializer) + for temporary in temporaries.values() + if temporary.initializer is not None] + lattice_size = (degree + 1) ** 3 + assert max(literal.size for literal in literals) <= lattice_size + assert sum(literal.size for literal in literals) < 10 * lattice_size + + +def test_bernstein_laplacian_bilinear_compact_codegen(): + import islpy as isl + import loopy as lp + + degree = 10 + collapsed = simplex_laplacian( + tetrahedron, "Bernstein", degree, scheme="collapsed") + canonical = simplex_laplacian( + tetrahedron, "Bernstein", degree, scheme="canonical") + collapsed_kernel, = compile_form( + collapsed, parameters=dict(mode="spectral")) + canonical_kernel, = compile_form( + canonical, parameters=dict(mode="spectral")) + + # At this degree the lower asymptotic complexity of sum factorisation + # outweighs its setup cost. + assert collapsed_kernel.flop_count < canonical_kernel.flop_count + + entrypoint = collapsed_kernel.ast.default_entrypoint + assert len(entrypoint.instructions) < 250 + assert len(entrypoint.temporary_variables) < 250 + code = lp.generate_code_v2(collapsed_kernel.ast).device_code() + assert sum(line.lstrip().startswith("for (") + for line in code.splitlines()) < 150 + + # A tetrahedral lattice has a loop whose bound depends on two parents. + assert max(domain.dim(isl.dim_type.param) + for domain in entrypoint.domains) >= 2 + + +@pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"]) +@pytest.mark.parametrize(('cell', 'order'), [(triangle, 5), (tetrahedron, 7)]) +def test_simplex_mass_bilinear(cell, family, order): + degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8)) + flops = [count_flops(simplex_mass(cell, family, degree)) + for degree in degrees] + rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees)) + assert (rates < order).all() + + +@pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"]) +@pytest.mark.parametrize(('cell', 'order'), [(triangle, 5), (tetrahedron, 7)]) +def test_simplex_laplacian_bilinear(cell, family, order): + degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8)) + flops = [count_flops(simplex_laplacian(cell, family, degree)) + for degree in degrees] + rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees)) + assert (rates < order).all() + + if __name__ == "__main__": import os import sys diff --git a/tsfc/kernel_interface/common.py b/tsfc/kernel_interface/common.py index 5d61a916aa..2aeddc5dc4 100644 --- a/tsfc/kernel_interface/common.py +++ b/tsfc/kernel_interface/common.py @@ -14,9 +14,10 @@ import numpy from FIAT.reference_element import TensorProductCell from finat.cell_tools import max_complex +from finat.duffy import DuffyElement from finat.quadrature import AbstractQuadratureRule from gem.node import traversal -from gem.optimise import constant_fold_zero +from gem.optimise import constant_fold_zero, unflatten_returns from gem.optimise import remove_componenttensors as prune from numpy import asarray from tsfc import fem @@ -210,6 +211,8 @@ def compile_gem(self, ctx): assignments.extend(mode.flatten(var_reps.items(), ctx['index_cache'])) if assignments: + # Rewrite flat FlattenedTensor scatters as jagged lattice loops + assignments = unflatten_returns(assignments) return_variables, expressions = zip(*assignments) else: return_variables = [] @@ -344,6 +347,11 @@ def set_quad_rule(params, cell, integral_type, functions): scheme = quad_rule fiat_cell = as_fiat_cell(cell) finat_elements = set(create_element(e) for e in elements if e.family() != "Real") + if (scheme == "default" and integral_type == "cell" + and any(isinstance(finat_el, DuffyElement) + for finat_el in finat_elements)): + # Duffy tabulation requires a collapsed-coordinate point set. + scheme = "collapsed" fiat_cells = [fiat_cell] + [finat_el.complex for finat_el in finat_elements] if any(c.is_macrocell() for c in fiat_cells): if len(set(c.get_spatial_dimension() for c in fiat_cells)) > 1: From fb36c55caba18151180cc37a16f1cd23e1daa404 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 6 Aug 2026 22:43:39 +0100 Subject: [PATCH 12/21] Benchmark Bernstein sum factorisation --- benchmarks/bernstein_laplacian.py | 112 ++++++++++++++++++ tests/tsfc/test_sum_factorisation.py | 167 +++++++++++++++++---------- 2 files changed, 218 insertions(+), 61 deletions(-) create mode 100755 benchmarks/bernstein_laplacian.py diff --git a/benchmarks/bernstein_laplacian.py b/benchmarks/bernstein_laplacian.py new file mode 100755 index 0000000000..312db8ecd3 --- /dev/null +++ b/benchmarks/bernstein_laplacian.py @@ -0,0 +1,112 @@ +#!/usr/bin/env python +"""Measure Bernstein Laplacian code generation on simplices.""" + +import argparse +import time + +import numpy + +from finat.ufl import FiniteElement, VectorElement +from tsfc import compile_form +from ufl import FunctionSpace, Mesh, TestFunction, TrialFunction, dx, grad, inner +from ufl.cell import Cell + + +def compile_target( + cell: Cell, degree: int, scheme: str) -> tuple[object, float]: + """Compile a Bernstein Laplacian bilinear form. + + Parameters + ---------- + cell + Reference simplex. + degree + Polynomial degree. + scheme + Quadrature scheme. + + Returns + ------- + kernel + Compiled TSFC kernel. + elapsed + Compilation time in seconds. + """ + mesh = Mesh(VectorElement("CG", cell, 1)) + space = FunctionSpace(mesh, FiniteElement("Bernstein", cell, degree)) + u = TrialFunction(space) + v = TestFunction(space) + form = inner(grad(u), grad(v)) * dx(scheme=scheme) + start = time.perf_counter() + kernel, = compile_form(form, parameters={"mode": "spectral"}) + return kernel, time.perf_counter() - start + + +def temporary_metrics(kernel: object) -> tuple[int, int, int, int, int]: + """Measure statically allocated Loopy temporaries. + + Parameters + ---------- + kernel + Compiled TSFC kernel. + + Returns + ------- + scalar_count + Number of scalar temporaries. + array_count + Number of array temporaries. + stored_values + Total scalar and array entries. + largest_array + Entries in the largest array temporary. + maximum_rank + Largest temporary tensor rank. + """ + temporaries = kernel.ast.default_entrypoint.temporary_variables.values() + shapes = [temporary.shape for temporary in temporaries] + array_sizes = [int(numpy.prod(shape)) for shape in shapes if shape] + return ( + sum(not shape for shape in shapes), + len(array_sizes), + sum(array_sizes) + sum(not shape for shape in shapes), + max(array_sizes, default=0), + max(map(len, shapes), default=0), + ) + + +def main() -> None: + """Print compiler metrics as copyable Markdown.""" + parser = argparse.ArgumentParser() + parser.add_argument( + "--cell", choices=("triangle", "tetrahedron"), + default="tetrahedron") + parser.add_argument("--degrees", nargs="+", type=int, default=(10,)) + parser.add_argument( + "--schemes", nargs="+", choices=("collapsed", "canonical"), + default=("collapsed", "canonical")) + args = parser.parse_args() + cell = Cell(args.cell) + + print("") + print("| cell | degree | scheme | compile (s) | flops | scalar temps | " + "array temps | stored values | bytes | largest | max rank | " + "AST lines |") + print("| :--- | ---: | :--- | ---: | ---: | ---: | ---: | ---: | " + "---: | ---: | ---: | ---: |") + for degree in args.degrees: + for scheme in args.schemes: + kernel, elapsed = compile_target(cell, degree, scheme) + source = str(kernel.ast) + nscalar, narray, nstored, largest, max_rank = \ + temporary_metrics(kernel) + print( + f"| {args.cell} | {degree} | {scheme} | {elapsed:.6f} | " + f"{kernel.flop_count:.0f} | {nscalar} | {narray} | " + f"{nstored} | {8 * nstored} | {largest} | {max_rank} | " + f"{len(source.splitlines())} |" + ) + + +if __name__ == "__main__": + main() diff --git a/tests/tsfc/test_sum_factorisation.py b/tests/tsfc/test_sum_factorisation.py index 1bf92871ba..d36fd78b49 100644 --- a/tests/tsfc/test_sum_factorisation.py +++ b/tests/tsfc/test_sum_factorisation.py @@ -257,111 +257,156 @@ def test_shared_map_is_tabulated_in_one_loop(cell, degree, expanded): assert selected <= count_loops(form) -def simplex_mass(cell, family, degree, scheme='collapsed'): - m = Mesh(VectorElement('CG', cell, 1)) - variant = None if family == "Bernstein" else "integral" - V = FunctionSpace(m, FiniteElement(family, cell, degree, variant=variant)) - u = TrialFunction(V) - v = TestFunction(V) +def bernstein_mass( + cell: object, degree: int, scheme: str = "collapsed") -> object: + """Construct a Bernstein mass form. + + Parameters + ---------- + cell + Reference simplex. + degree + Polynomial degree. + scheme + Quadrature scheme. + + Returns + ------- + object + UFL bilinear form. + """ + mesh = Mesh(VectorElement("CG", cell, 1)) + space = FunctionSpace(mesh, FiniteElement("Bernstein", cell, degree)) + u = TrialFunction(space) + v = TestFunction(space) return inner(u, v) * dx(scheme=scheme) -def simplex_laplacian(cell, family, degree, scheme='collapsed'): - m = Mesh(VectorElement('CG', cell, 1)) - variant = None if family == "Bernstein" else "integral" - V = FunctionSpace(m, FiniteElement(family, cell, degree, variant=variant)) - u = TrialFunction(V) - v = TestFunction(V) +def bernstein_laplacian( + cell: object, degree: int, scheme: str = "collapsed") -> object: + """Construct a Bernstein Laplacian form. + + Parameters + ---------- + cell + Reference simplex. + degree + Polynomial degree. + scheme + Quadrature scheme. + + Returns + ------- + object + UFL bilinear form. + """ + mesh = Mesh(VectorElement("CG", cell, 1)) + space = FunctionSpace(mesh, FiniteElement("Bernstein", cell, degree)) + u = TrialFunction(space) + v = TestFunction(space) return inner(grad(u), grad(v)) * dx(scheme=scheme) -@pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"]) -@pytest.mark.parametrize(('cell', 'order'), [(triangle, 3), (tetrahedron, 4)]) -def test_simplex_mass_action(cell, family, order): +@pytest.mark.parametrize(("cell", "order"), [(triangle, 3), (tetrahedron, 4)]) +def test_bernstein_mass_action(cell: object, order: float) -> None: degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8)) - flops = [count_flops(action(simplex_mass(cell, family, degree))) - for degree in degrees] + flops = [ + count_flops(action(bernstein_mass(cell, degree))) + for degree in degrees + ] rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees)) assert (rates < order).all() -@pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"]) -@pytest.mark.parametrize(('cell', 'order'), [(triangle, 3), (tetrahedron, 4.4)]) -def test_simplex_laplacian_action(cell, family, order): +@pytest.mark.parametrize( + ("cell", "order"), [(triangle, 3), (tetrahedron, 4.4)]) +def test_bernstein_laplacian_action(cell: object, order: float) -> None: degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8)) - flops = [count_flops(action(simplex_laplacian(cell, family, degree))) - for degree in degrees] + flops = [ + count_flops(action(bernstein_laplacian(cell, degree))) + for degree in degrees + ] rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees)) assert (rates < order).all() -@pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"]) -def test_simplex_laplacian_action_compact_codegen(family): - form = action(simplex_laplacian(triangle, family, 3)) - kernel, = compile_form(form, parameters=dict(mode='spectral')) - temporaries = kernel.ast.default_entrypoint.temporary_variables - assert len(temporaries) < 100 - - -def test_bernstein_laplacian_action_compact_literals(): +def test_bernstein_laplacian_action_compact_literals() -> None: degree = 5 - form = action(simplex_laplacian(tetrahedron, "Bernstein", degree)) - kernel, = compile_form(form, parameters=dict(mode='spectral')) + form = action(bernstein_laplacian(tetrahedron, degree)) + kernel, = compile_form(form, parameters={"mode": "spectral"}) temporaries = kernel.ast.default_entrypoint.temporary_variables - literals = [numpy.asarray(temporary.initializer) - for temporary in temporaries.values() - if temporary.initializer is not None] + literals = [ + numpy.asarray(temporary.initializer) + for temporary in temporaries.values() + if temporary.initializer is not None + ] lattice_size = (degree + 1) ** 3 assert max(literal.size for literal in literals) <= lattice_size assert sum(literal.size for literal in literals) < 10 * lattice_size -def test_bernstein_laplacian_bilinear_compact_codegen(): +def test_bernstein_laplacian_bilinear_compact_codegen() -> None: import islpy as isl import loopy as lp degree = 10 - collapsed = simplex_laplacian( - tetrahedron, "Bernstein", degree, scheme="collapsed") - canonical = simplex_laplacian( - tetrahedron, "Bernstein", degree, scheme="canonical") + collapsed = bernstein_laplacian( + tetrahedron, degree, scheme="collapsed") + canonical = bernstein_laplacian( + tetrahedron, degree, scheme="canonical") collapsed_kernel, = compile_form( - collapsed, parameters=dict(mode="spectral")) + collapsed, parameters={"mode": "spectral"}) canonical_kernel, = compile_form( - canonical, parameters=dict(mode="spectral")) + canonical, parameters={"mode": "spectral"}) - # At this degree the lower asymptotic complexity of sum factorisation - # outweighs its setup cost. assert collapsed_kernel.flop_count < canonical_kernel.flop_count entrypoint = collapsed_kernel.ast.default_entrypoint - assert len(entrypoint.instructions) < 250 - assert len(entrypoint.temporary_variables) < 250 + collapsed_shapes = [ + temporary.shape + for temporary in entrypoint.temporary_variables.values() + ] + canonical_shapes = [ + temporary.shape + for temporary in + canonical_kernel.ast.default_entrypoint.temporary_variables.values() + ] + assert max(map(len, collapsed_shapes)) <= 5 + assert sum(map(numpy.prod, collapsed_shapes)) \ + < sum(map(numpy.prod, canonical_shapes)) + code = lp.generate_code_v2(collapsed_kernel.ast).device_code() - assert sum(line.lstrip().startswith("for (") - for line in code.splitlines()) < 150 + assert sum( + line.lstrip().startswith("for (") + for line in code.splitlines() + ) < 150 # A tetrahedral lattice has a loop whose bound depends on two parents. - assert max(domain.dim(isl.dim_type.param) - for domain in entrypoint.domains) >= 2 + assert max( + domain.dim(isl.dim_type.param) + for domain in entrypoint.domains + ) >= 2 -@pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"]) -@pytest.mark.parametrize(('cell', 'order'), [(triangle, 5), (tetrahedron, 7)]) -def test_simplex_mass_bilinear(cell, family, order): +@pytest.mark.parametrize(("cell", "order"), [(triangle, 5), (tetrahedron, 7)]) +def test_bernstein_mass_bilinear(cell: object, order: float) -> None: degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8)) - flops = [count_flops(simplex_mass(cell, family, degree)) - for degree in degrees] + flops = [ + count_flops(bernstein_mass(cell, degree)) + for degree in degrees + ] rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees)) assert (rates < order).all() -@pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"]) -@pytest.mark.parametrize(('cell', 'order'), [(triangle, 5), (tetrahedron, 7)]) -def test_simplex_laplacian_bilinear(cell, family, order): +@pytest.mark.parametrize(("cell", "order"), [(triangle, 5), (tetrahedron, 7)]) +def test_bernstein_laplacian_bilinear( + cell: object, order: float) -> None: degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8)) - flops = [count_flops(simplex_laplacian(cell, family, degree)) - for degree in degrees] + flops = [ + count_flops(bernstein_laplacian(cell, degree)) + for degree in degrees + ] rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees)) assert (rates < order).all() From d528ba3069fd3201e3f13406cd9399b7ca6bb205 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 14 Aug 2026 11:52:41 +0100 Subject: [PATCH 13/21] 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 68f0fc7f28..1aaface191 100644 --- a/tsfc/loopy.py +++ b/tsfc/loopy.py @@ -126,7 +126,6 @@ def __init__(self, target=None): self.tabulated = (None, ()) # axes and inames of the preceding tabulation self.index_parents = {} # iname -> parent inames bounding a jagged index self.gem_to_pymbolic = {} # gem node -> pymbolic variable - self.compact_indices = {} # temporary -> compact index layout self.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 b946096427f0dd566020e8088b585bb4f0b2a8a4 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 28 Aug 2026 21:45:20 +0100 Subject: [PATCH 14/21] Bound simplex code generation storage --- benchmarks/bernstein_laplacian.py | 108 +++++++++++++++++++++++++-- tests/tsfc/test_sum_factorisation.py | 20 ++--- tsfc/kernel_interface/common.py | 27 ++++++- tsfc/loopy.py | 78 ++++++++++++++++++- 4 files changed, 213 insertions(+), 20 deletions(-) diff --git a/benchmarks/bernstein_laplacian.py b/benchmarks/bernstein_laplacian.py index 312db8ecd3..d1820269c3 100755 --- a/benchmarks/bernstein_laplacian.py +++ b/benchmarks/bernstein_laplacian.py @@ -2,8 +2,17 @@ """Measure Bernstein Laplacian code generation on simplices.""" import argparse +import ctypes +from math import comb +import os +from pathlib import Path +import shlex +import statistics +import subprocess +import tempfile import time +import loopy as lp import numpy from finat.ufl import FiniteElement, VectorElement @@ -63,18 +72,86 @@ def temporary_metrics(kernel: object) -> tuple[int, int, int, int, int]: maximum_rank Largest temporary tensor rank. """ - temporaries = kernel.ast.default_entrypoint.temporary_variables.values() + temporaries = tuple( + kernel.ast.default_entrypoint.temporary_variables.values()) shapes = [temporary.shape for temporary in temporaries] + storage_shapes = [temporary.shape for temporary in temporaries + if temporary.base_storage is None] array_sizes = [int(numpy.prod(shape)) for shape in shapes if shape] + storage_sizes = [int(numpy.prod(shape)) if shape else 1 + for shape in storage_shapes] return ( sum(not shape for shape in shapes), len(array_sizes), - sum(array_sizes) + sum(not shape for shape in shapes), + sum(storage_sizes), max(array_sizes, default=0), max(map(len, shapes), default=0), ) +def direct_runtime( + kernel: object, cell: Cell, degree: int, calls: int, + repeats: int) -> tuple[numpy.ndarray, float]: + """Compile and time direct calls to a generated C kernel. + + Parameters + ---------- + kernel + Compiled TSFC kernel. + cell + Reference simplex. + degree + Bernstein polynomial degree. + calls + Kernel calls in each timed sample. + repeats + Number of timed samples. + + Returns + ------- + output + Tensor produced by one kernel call. + elapsed + Median seconds per direct kernel call. + """ + dimension = cell.topological_dimension + ndofs = comb(degree + dimension, dimension) + coordinates = numpy.vstack(( + numpy.zeros((1, dimension)), numpy.eye(dimension) + )).astype(numpy.float64) + output = numpy.zeros((ndofs, ndofs), dtype=numpy.float64) + source = lp.generate_code_v2(kernel.ast).device_code() + + with tempfile.TemporaryDirectory(prefix="bernstein-kernel-") as directory: + directory = Path(directory) + source_path = directory / "kernel.c" + library_path = directory / "kernel.so" + source_path.write_text(source) + compiler = shlex.split(os.environ.get("CC", "cc")) + subprocess.run( + [*compiler, "-O3", "-march=native", "-shared", "-fPIC", + str(source_path), "-lm", "-o", str(library_path)], + check=True, + ) + library = ctypes.CDLL(str(library_path)) + function = getattr(library, kernel.ast.default_entrypoint.name) + pointer = ctypes.POINTER(ctypes.c_double) + function.argtypes = (pointer, pointer) + output_pointer = output.ctypes.data_as(pointer) + coordinate_pointer = coordinates.ctypes.data_as(pointer) + + output.fill(0) + function(output_pointer, coordinate_pointer) + reference = output.copy() + samples = [] + for _ in range(repeats): + start = time.perf_counter() + for _ in range(calls): + function(output_pointer, coordinate_pointer) + samples.append((time.perf_counter() - start) / calls) + return reference, statistics.median(samples) + + def main() -> None: """Print compiler metrics as copyable Markdown.""" parser = argparse.ArgumentParser() @@ -85,26 +162,43 @@ def main() -> None: parser.add_argument( "--schemes", nargs="+", choices=("collapsed", "canonical"), default=("collapsed", "canonical")) + parser.add_argument("--runtime-calls", type=int, default=5) + parser.add_argument("--runtime-repeats", type=int, default=3) args = parser.parse_args() cell = Cell(args.cell) print("") - print("| cell | degree | scheme | compile (s) | flops | scalar temps | " - "array temps | stored values | bytes | largest | max rank | " + print("| cell | degree | scheme | compile (s) | runtime (ms/call) | " + "max error | flops | scalar temps | " + "array temps | allocated values | bytes | largest | max rank | " "AST lines |") print("| :--- | ---: | :--- | ---: | ---: | ---: | ---: | ---: | " - "---: | ---: | ---: | ---: |") + "---: | ---: | ---: | ---: | ---: | ---: |") for degree in args.degrees: + rows = [] for scheme in args.schemes: kernel, elapsed = compile_target(cell, degree, scheme) source = str(kernel.ast) nscalar, narray, nstored, largest, max_rank = \ temporary_metrics(kernel) + output, runtime = direct_runtime( + kernel, cell, degree, args.runtime_calls, + args.runtime_repeats) + rows.append((scheme, output, runtime, elapsed, kernel.flop_count, + nscalar, narray, nstored, largest, max_rank, + len(source.splitlines()))) + reference = next((output for scheme, output, *_ in rows + if scheme == "canonical"), None) + for (scheme, output, runtime, elapsed, flops, nscalar, narray, + nstored, largest, max_rank, ast_lines) in rows: + error = (numpy.max(numpy.abs(output - reference)) + if reference is not None else numpy.nan) print( f"| {args.cell} | {degree} | {scheme} | {elapsed:.6f} | " - f"{kernel.flop_count:.0f} | {nscalar} | {narray} | " + f"{1000 * runtime:.6f} | {error:.3e} | {flops:.0f} | " + f"{nscalar} | {narray} | " f"{nstored} | {8 * nstored} | {largest} | {max_rank} | " - f"{len(source.splitlines())} |" + f"{ast_lines} |" ) diff --git a/tests/tsfc/test_sum_factorisation.py b/tests/tsfc/test_sum_factorisation.py index d36fd78b49..8ae515c646 100644 --- a/tests/tsfc/test_sum_factorisation.py +++ b/tests/tsfc/test_sum_factorisation.py @@ -362,24 +362,26 @@ def test_bernstein_laplacian_bilinear_compact_codegen() -> None: assert collapsed_kernel.flop_count < canonical_kernel.flop_count entrypoint = collapsed_kernel.ast.default_entrypoint + collapsed_temporaries = tuple(entrypoint.temporary_variables.values()) + canonical_temporaries = tuple( + canonical_kernel.ast.default_entrypoint.temporary_variables.values()) collapsed_shapes = [ temporary.shape - for temporary in entrypoint.temporary_variables.values() - ] - canonical_shapes = [ - temporary.shape - for temporary in - canonical_kernel.ast.default_entrypoint.temporary_variables.values() + for temporary in collapsed_temporaries ] assert max(map(len, collapsed_shapes)) <= 5 - assert sum(map(numpy.prod, collapsed_shapes)) \ - < sum(map(numpy.prod, canonical_shapes)) + assert sum(numpy.prod(temporary.shape) + for temporary in collapsed_temporaries + if temporary.base_storage is None) \ + < sum(numpy.prod(temporary.shape) + for temporary in canonical_temporaries + if temporary.base_storage is None) code = lp.generate_code_v2(collapsed_kernel.ast).device_code() assert sum( line.lstrip().startswith("for (") for line in code.splitlines() - ) < 150 + ) < 500 # A tetrahedral lattice has a loop whose bound depends on two parents. assert max( diff --git a/tsfc/kernel_interface/common.py b/tsfc/kernel_interface/common.py index 2aeddc5dc4..ecab498dc9 100644 --- a/tsfc/kernel_interface/common.py +++ b/tsfc/kernel_interface/common.py @@ -27,6 +27,25 @@ from tsfc.logging import logger +def _has_product_lattice_scatter(variable: gem.Node) -> bool: + """Check whether a return scatters a product of jagged lattices.""" + if isinstance(variable, gem.Indexed): + indices = variable.multiindex + elif isinstance(variable, gem.FlexiblyIndexed): + indices = tuple( + index + for _, dimension in variable.dim2idxs + for index, _ in dimension + ) + else: + return False + return ( + any(isinstance(index, gem.VariableIndex) for index in indices) + and sum(isinstance(index, gem.JaggedIndex) and not index.parents + for index in variable.free_indices) > 1 + ) + + class KernelBuilderBase(KernelInterface): """Helper class for building local assembly kernels.""" @@ -235,8 +254,14 @@ def compile_gem(self, ctx): # Construct ImperoC assignments = list(zip(return_variables, expressions)) index_ordering = get_index_ordering(ctx['quadrature_indices'], return_variables) + assignment_group_size = ( + 8 if any(map(_has_product_lattice_scatter, return_variables)) + else None + ) try: - impero_c = impero_utils.compile_gem(assignments, index_ordering, remove_zeros=True) + impero_c = impero_utils.compile_gem( + assignments, index_ordering, remove_zeros=True, + assignment_group_size=assignment_group_size) except impero_utils.NoopError: impero_c = None return impero_c, oriented, needs_cell_sizes, tabulations, active_variables diff --git a/tsfc/loopy.py b/tsfc/loopy.py index 1aaface191..ebc9ff9cf6 100644 --- a/tsfc/loopy.py +++ b/tsfc/loopy.py @@ -8,6 +8,7 @@ from collections import defaultdict, OrderedDict from gem import gem, impero as imp +from gem.impero_utils import temp_refcount from gem.node import Memoizer import islpy as isl @@ -126,6 +127,7 @@ def __init__(self, target=None): self.tabulated = (None, ()) # axes and inames of the preceding tabulation self.index_parents = {} # iname -> parent inames bounding a jagged index self.gem_to_pymbolic = {} # gem node -> pymbolic variable + self.compact_indices = {} # temporary -> compact index layout self.name_gen = UniqueNameGenerator() self.target = target self.loop_priorities = set() # used to avoid disadvantageous loop interchanges @@ -175,6 +177,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: @@ -214,6 +223,55 @@ def active_indices(mapping, ctx): ctx.active_indices.pop(key) +def _temporary_base_storage(impero_c, descriptors): + """Alias equal-shaped temporaries with disjoint Impero lifetimes.""" + numbering = {temporary: temporary for temporary in impero_c.temporaries} + uses = defaultdict(list) + position = 0 + + def visit(node): + nonlocal position + if isinstance(node, imp.Terminal): + for temporary in temp_refcount(numbering, node): + uses[temporary].append(position) + position += 1 + return + for child in node.children: + visit(child) + + visit(impero_c.tree) + intervals = { + temporary: (min(positions), max(positions)) + for temporary, positions in uses.items() + } + + def overlap(left, right): + return not (left[1] < right[0] or right[1] < left[0]) + + pools = defaultdict(list) + storage = {} + for temporary, dtype, shape, _ in descriptors: + interval = intervals.get(temporary) + if isinstance(temporary, gem.Constant) or interval is None or not shape: + continue + key = str(dtype), shape + for name, occupied in pools[key]: + if not any(overlap(interval, other) for other in occupied): + occupied.append(interval) + storage[temporary] = name + break + else: + name = f"storage{sum(map(len, pools.values()))}" + pools[key].append((name, [interval])) + storage[temporary] = name + + counts = defaultdict(int) + for name in storage.values(): + counts[name] += 1 + return {temporary: name for temporary, name in storage.items() + if counts[name] > 1} + + def generate(impero_c, args, scalar_type, kernel_name="loopy_kernel", index_names=[], return_increments=True, log=False): """Generates loopy code. @@ -236,13 +294,26 @@ def generate(impero_c, args, scalar_type, kernel_name="loopy_kernel", index_name # Create arguments data = list(args) - for i, (temp, dtype) in enumerate(assign_dtypes(impero_c.temporaries, scalar_type)): + descriptors = [] + for temp, dtype in assign_dtypes(impero_c.temporaries, scalar_type): + if isinstance(temp, gem.Constant): + shape, layout = temp.shape, None + else: + shape, layout = gem.compact_index_layout( + tuple(ctx.indices[temp])) + shape += temp.shape + descriptors.append((temp, dtype, shape, layout)) + base_storage = _temporary_base_storage(impero_c, descriptors) + for i, (temp, dtype, shape, layout) in enumerate(descriptors): name = "t%d" % i if isinstance(temp, gem.Constant): data.append(lp.TemporaryVariable(name, shape=temp.shape, dtype=dtype, initializer=temp.array, address_space=lp.AddressSpace.LOCAL, read_only=True)) else: - shape = tuple([i.extent for i in ctx.indices[temp]]) + temp.shape - data.append(lp.TemporaryVariable(name, shape=shape, dtype=dtype, initializer=None, address_space=lp.AddressSpace.LOCAL, read_only=False)) + data.append(lp.TemporaryVariable( + name, shape=shape, dtype=dtype, initializer=None, + address_space=lp.AddressSpace.LOCAL, read_only=False, + base_storage=base_storage.get(temp))) + ctx.compact_indices[temp] = layout ctx.gem_to_pymbolic[temp] = p.Variable(name) # Create instructions @@ -274,6 +345,7 @@ def generate(impero_c, args, scalar_type, kernel_name="loopy_kernel", index_name preambles=preamble, loop_priority=frozenset(ctx.loop_priorities), ) + knl = lp.allocate_temporaries_for_base_storage(knl) return knl, event_name From 7044d3c36a03fc3975d80bb91c116d5fe46f85e6 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 29 Aug 2026 00:54:35 +0100 Subject: [PATCH 15/21] Schedule terminal reductions inside the argument loops get_index_ordering put every quadrature index outermost. Impero derives every loop nest from that one global order, so a reduction ordered outside its own free indices has to accumulate into a temporary carrying all of them. Where the argument lattice is factorised those free indices are the whole output, so the degree-10 tetrahedral mass bilinear staged the entire (286, 286) element matrix and then traversed both lattices a second time only to combine and scatter it. Quadrature-outermost is right in general: it is what lets each stage of a tensor-product contraction shed its quadrature axis, and flipping it globally regresses test_contraction_storage_rate on quadrilaterals and tensor-product cells from O(p) to O(p^2) stages. So offer both orderings and keep the cheaper. _terminal_reductions finds quadrature indices whose reduction spans the whole output and is not the root of its assignment; a root reduction is already absorbed by ReturnAccumulate, a buried one is not. index_orderings returns the default plus one that moves exactly those indices innermost, and the candidate with less declared temporary storage wins. The two orderings have identical flop counts, so the choice is only about where the intermediates live. Peak live storage would be the better metric in principle, but the two metrics agree on every case measured -- 32 two-candidate choices across mass and Laplacian, bilinear and action, triangles and tetrahedra at degrees 3 to 10 -- and the margin between orderings is two orders of magnitude wider than the margin between metrics. With the output-shaped temporary gone, two workarounds are no longer needed: the batch size of eight in _has_product_lattice_scatter, and the lifetime aliasing in _temporary_base_storage, which was saving 583 words out of 157756. The loop-count assertion goes back to the original < 150; the degree-10 Laplacian emits 136 loops, against 254 with aliasing. Degree 10 on tetrahedra, declared temporary storage in words: mass bilinear 87241 to 7997, Laplacian bilinear 3089666 to 138204. The largest staged matrix is (66, 66), a product of two 2-simplex lattices. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_013NoTXWyj2fVdJTHnMDFB4k --- tests/tsfc/test_sum_factorisation.py | 42 ++++++++++-- tsfc/kernel_interface/common.py | 99 +++++++++++++++++++--------- tsfc/loopy.py | 93 +++++++++++--------------- 3 files changed, 141 insertions(+), 93 deletions(-) diff --git a/tests/tsfc/test_sum_factorisation.py b/tests/tsfc/test_sum_factorisation.py index 8ae515c646..4d6d3b8064 100644 --- a/tests/tsfc/test_sum_factorisation.py +++ b/tests/tsfc/test_sum_factorisation.py @@ -1,3 +1,5 @@ +from math import comb + import numpy import pytest @@ -370,18 +372,15 @@ def test_bernstein_laplacian_bilinear_compact_codegen() -> None: for temporary in collapsed_temporaries ] assert max(map(len, collapsed_shapes)) <= 5 - assert sum(numpy.prod(temporary.shape) - for temporary in collapsed_temporaries - if temporary.base_storage is None) \ + assert sum(map(numpy.prod, collapsed_shapes)) \ < sum(numpy.prod(temporary.shape) - for temporary in canonical_temporaries - if temporary.base_storage is None) + for temporary in canonical_temporaries) code = lp.generate_code_v2(collapsed_kernel.ast).device_code() assert sum( line.lstrip().startswith("for (") for line in code.splitlines() - ) < 500 + ) < 150 # A tetrahedral lattice has a loop whose bound depends on two parents. assert max( @@ -390,6 +389,37 @@ def test_bernstein_laplacian_bilinear_compact_codegen() -> None: ) >= 2 +def test_bernstein_bilinear_contracts_into_the_output() -> None: + # The last one-dimensional contraction runs inside both argument + # lattices and accumulates into the scattered output. Nothing between + # the contraction and the output holds a whole element matrix, so the + # dominant loop nest is two simplex lattices plus one quadrature loop. + degree = 10 + dimension = 3 + kernel, = compile_form( + bernstein_mass(tetrahedron, degree, scheme="collapsed"), + parameters={"mode": "spectral"}) + entrypoint = kernel.ast.default_entrypoint + + nodes = comb(degree + dimension, dimension) + writable = [temporary + for temporary in entrypoint.temporary_variables.values() + if temporary.initializer is None] + assert max(numpy.prod(temporary.shape, dtype=int) + for temporary in writable) < nodes ** 2 + + scatters = [instruction for instruction in entrypoint.instructions + if "A" in instruction.write_dependency_names()] + assert scatters + scatter_inames, = {frozenset(instruction.within_inames) + for instruction in scatters} + arguments = 2 + assert len(scatter_inames) == arguments * dimension + # The quadrature reduction is nested inside that scatter. + assert any(instruction.within_inames > scatter_inames + for instruction in entrypoint.instructions) + + @pytest.mark.parametrize(("cell", "order"), [(triangle, 5), (tetrahedron, 7)]) def test_bernstein_mass_bilinear(cell: object, order: float) -> None: degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8)) diff --git a/tsfc/kernel_interface/common.py b/tsfc/kernel_interface/common.py index ecab498dc9..cb50860dde 100644 --- a/tsfc/kernel_interface/common.py +++ b/tsfc/kernel_interface/common.py @@ -9,6 +9,7 @@ from ufl.domain import extract_unique_domain import gem +import gem.gem import gem.impero_utils as impero_utils import petsctools import numpy @@ -27,25 +28,6 @@ from tsfc.logging import logger -def _has_product_lattice_scatter(variable: gem.Node) -> bool: - """Check whether a return scatters a product of jagged lattices.""" - if isinstance(variable, gem.Indexed): - indices = variable.multiindex - elif isinstance(variable, gem.FlexiblyIndexed): - indices = tuple( - index - for _, dimension in variable.dim2idxs - for index, _ in dimension - ) - else: - return False - return ( - any(isinstance(index, gem.VariableIndex) for index in indices) - and sum(isinstance(index, gem.JaggedIndex) and not index.parents - for index in variable.free_indices) > 1 - ) - - class KernelBuilderBase(KernelInterface): """Helper class for building local assembly kernels.""" @@ -253,17 +235,18 @@ def compile_gem(self, ctx): active_variables = gem.extract_type(expressions, gem.Variable) # Construct ImperoC assignments = list(zip(return_variables, expressions)) - index_ordering = get_index_ordering(ctx['quadrature_indices'], return_variables) - assignment_group_size = ( - 8 if any(map(_has_product_lattice_scatter, return_variables)) - else None - ) - try: - impero_c = impero_utils.compile_gem( - assignments, index_ordering, remove_zeros=True, - assignment_group_size=assignment_group_size) - except impero_utils.NoopError: + candidates = [] + for index_ordering in index_orderings( + ctx['quadrature_indices'], return_variables, assignments): + try: + candidates.append(impero_utils.compile_gem( + assignments, index_ordering, remove_zeros=True)) + except impero_utils.NoopError: + candidates.append(None) + if any(candidate is None for candidate in candidates): impero_c = None + else: + impero_c = min(candidates, key=_storage_cost) return impero_c, oriented, needs_cell_sizes, tabulations, active_variables def fem_config(self): @@ -393,10 +376,64 @@ def set_quad_rule(params, cell, integral_type, functions): type(quad_rule)) -def get_index_ordering(quadrature_indices, return_variables): +def _spans_output(node, output_indices): + """Check whether a node reduces a summand that spans the whole output.""" + return (isinstance(node, gem.IndexSum) + and set(node.free_indices) == output_indices) + + +def _terminal_reductions(assignments): + """Quadrature indices whose reduction cannot shrink what it accumulates. + + Such a reduction has no tensor smaller than the output to accumulate + into, so running its loop outside the argument loops costs an + output-shaped temporary. Placing it innermost instead makes the + accumulator a scalar, at the cost of giving every stage that feeds it + a quadrature axis, so neither placement dominates and both are costed. + """ + terminal = set() + for variable, expression in assignments: + output_indices = set(variable.free_indices) + if _spans_output(expression, output_indices): + # A root reduction accumulates straight into the output, so + # the outer placement costs nothing extra. + continue + for node in traversal((expression,)): + if _spans_output(node, output_indices): + terminal.update(node.multiindex) + return terminal + + +def index_orderings(quadrature_indices, return_variables, assignments=()): + """Return the candidate outermost loop orderings, best guess first. + + Quadrature loops run outside the argument loops, which keeps every + intermediate contraction stage free of a quadrature axis. When an + output-shaped reduction cannot accumulate into the output itself, + that placement also forces an output-shaped temporary, so the + ordering that runs it innermost is offered as an alternative. + """ split_argument_indices = tuple(chain(*(var.index_ordering() for var in return_variables))) - return tuple(quadrature_indices) + split_argument_indices + quadrature_indices = tuple(quadrature_indices) + default = quadrature_indices + split_argument_indices + terminal = _terminal_reductions(assignments) + if not terminal: + return (default,) + head = tuple(i for i in quadrature_indices if i not in terminal) + tail = tuple(i for i in quadrature_indices if i in terminal) + return (default, head + split_argument_indices + tail) + + +def _storage_cost(impero_c): + """Total declared size of an Impero program's temporaries.""" + total = 0 + for temporary in impero_c.temporaries: + if isinstance(temporary, gem.gem.Constant): + continue + shape, _ = gem.compact_index_layout(tuple(impero_c.indices[temporary])) + total += numpy.prod(shape + temporary.shape, dtype=int) + return total def get_index_names(quadrature_indices, argument_multiindices, index_cache): diff --git a/tsfc/loopy.py b/tsfc/loopy.py index ebc9ff9cf6..c71372c018 100644 --- a/tsfc/loopy.py +++ b/tsfc/loopy.py @@ -8,7 +8,6 @@ from collections import defaultdict, OrderedDict from gem import gem, impero as imp -from gem.impero_utils import temp_refcount from gem.node import Memoizer import islpy as isl @@ -128,6 +127,7 @@ def __init__(self, target=None): self.index_parents = {} # iname -> parent inames bounding a jagged index self.gem_to_pymbolic = {} # gem node -> pymbolic variable self.compact_indices = {} # temporary -> compact index layout + self.lattice_ranks = {} # lattice shape -> rank table and its entries self.name_gen = UniqueNameGenerator() self.target = target self.loop_priorities = set() # used to avoid disadvantageous loop interchanges @@ -179,8 +179,7 @@ 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) + self.lattice_rank(item) if isinstance(item, tuple) else self.active_indices[item] for item in self.compact_indices[node]) return p.Subscript(pym, indices) if indices else pym @@ -199,6 +198,38 @@ def _gem_to_pym_var(self, node): self.gem_to_pymbolic[node] = pym return pym + def lattice_rank(self, component): + """Look up the compact rank of the active simplex lattice point. + + A simplex lattice is stored along one compact dimension, so a + temporary indexed by the lattice is subscripted by the rank of + the point rather than by the lattice indices themselves. The + ranks are tabulated once per lattice shape, which keeps the + lattice visible in the AST and keeps the polynomial that defines + the rank out of the loop body. + """ + # Equal degree and dimension describe a simplex lattice fully, + # so every lattice of one shape shares a single table. + shape = (component[0].extent, len(component)) + try: + table, _ = self.lattice_ranks[shape] + except KeyError: + table = p.Variable(self.name_gen("lattice_rank")) + self.lattice_ranks[shape] = ( + table, gem.simplex_lattice_ranks(component)) + return p.Subscript(table, tuple(self.active_indices[index] + for index in component)) + + def lattice_rank_tables(self, address_space): + """Declare a read-only table for each tabulated simplex lattice.""" + return [ + lp.TemporaryVariable( + table.name, shape=ranks.shape, dtype=ranks.dtype, + initializer=ranks, address_space=address_space, + read_only=True) + for table, ranks in self.lattice_ranks.values() + ] + def active_inames(self): # Return all active indices return frozenset([i.name for i in self.active_indices.values()]) @@ -223,55 +254,6 @@ def active_indices(mapping, ctx): ctx.active_indices.pop(key) -def _temporary_base_storage(impero_c, descriptors): - """Alias equal-shaped temporaries with disjoint Impero lifetimes.""" - numbering = {temporary: temporary for temporary in impero_c.temporaries} - uses = defaultdict(list) - position = 0 - - def visit(node): - nonlocal position - if isinstance(node, imp.Terminal): - for temporary in temp_refcount(numbering, node): - uses[temporary].append(position) - position += 1 - return - for child in node.children: - visit(child) - - visit(impero_c.tree) - intervals = { - temporary: (min(positions), max(positions)) - for temporary, positions in uses.items() - } - - def overlap(left, right): - return not (left[1] < right[0] or right[1] < left[0]) - - pools = defaultdict(list) - storage = {} - for temporary, dtype, shape, _ in descriptors: - interval = intervals.get(temporary) - if isinstance(temporary, gem.Constant) or interval is None or not shape: - continue - key = str(dtype), shape - for name, occupied in pools[key]: - if not any(overlap(interval, other) for other in occupied): - occupied.append(interval) - storage[temporary] = name - break - else: - name = f"storage{sum(map(len, pools.values()))}" - pools[key].append((name, [interval])) - storage[temporary] = name - - counts = defaultdict(int) - for name in storage.values(): - counts[name] += 1 - return {temporary: name for temporary, name in storage.items() - if counts[name] > 1} - - def generate(impero_c, args, scalar_type, kernel_name="loopy_kernel", index_names=[], return_increments=True, log=False): """Generates loopy code. @@ -303,7 +285,6 @@ def generate(impero_c, args, scalar_type, kernel_name="loopy_kernel", index_name tuple(ctx.indices[temp])) shape += temp.shape descriptors.append((temp, dtype, shape, layout)) - base_storage = _temporary_base_storage(impero_c, descriptors) for i, (temp, dtype, shape, layout) in enumerate(descriptors): name = "t%d" % i if isinstance(temp, gem.Constant): @@ -311,14 +292,15 @@ def generate(impero_c, args, scalar_type, kernel_name="loopy_kernel", index_name else: data.append(lp.TemporaryVariable( name, shape=shape, dtype=dtype, initializer=None, - address_space=lp.AddressSpace.LOCAL, read_only=False, - base_storage=base_storage.get(temp))) + address_space=lp.AddressSpace.LOCAL, read_only=False)) ctx.compact_indices[temp] = layout ctx.gem_to_pymbolic[temp] = p.Variable(name) # Create instructions instructions = statement(impero_c.tree, ctx) + data.extend(ctx.lattice_rank_tables(lp.AddressSpace.LOCAL)) + # add a no-op touching all kernel arguments to make sure they # are not silently dropped noop = lp.CInstruction( @@ -345,7 +327,6 @@ def generate(impero_c, args, scalar_type, kernel_name="loopy_kernel", index_name preambles=preamble, loop_priority=frozenset(ctx.loop_priorities), ) - knl = lp.allocate_temporaries_for_base_storage(knl) return knl, event_name From c69ae811cc5aab3c85820b2f6d47f9df218cd555 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 29 Aug 2026 01:01:27 +0100 Subject: [PATCH 16/21] Keep get_index_ordering for single-ordering callers compile_expression_dual_evaluation compiles one ordering rather than costing several, so it needs the scalar form rather than the candidate list. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_013NoTXWyj2fVdJTHnMDFB4k --- tsfc/kernel_interface/common.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tsfc/kernel_interface/common.py b/tsfc/kernel_interface/common.py index cb50860dde..a807b3a161 100644 --- a/tsfc/kernel_interface/common.py +++ b/tsfc/kernel_interface/common.py @@ -404,6 +404,15 @@ def _terminal_reductions(assignments): return terminal +def get_index_ordering(quadrature_indices, return_variables, assignments=()): + """Return the single best-guess outermost loop ordering. + + For callers that compile one ordering rather than costing several. + """ + return index_orderings(quadrature_indices, return_variables, + assignments)[0] + + def index_orderings(quadrature_indices, return_variables, assignments=()): """Return the candidate outermost loop orderings, best guess first. From 34c5f58322b10a8c3e51621d1b8ab3ed66c271da Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 29 Aug 2026 02:08:06 +0100 Subject: [PATCH 17/21] DROP BEFORE MERGE: install the FIAT stack before anything imports Firedrake The install step ends with firedrake-clean, which imports Firedrake, and so imports the tsfc that needs the GEM changes in the FIAT stack. Swapping the stack in from a step after that one leaves firedrake-clean to run against the FIAT that pyproject.toml resolved from main: ImportError: cannot import name 'eliminate_deltas' from 'gem.optimise' The job never reached firedrake-check. Install the stack inline instead, after the Firedrake install and before firedrake-clean, which is what pbrubeck/form-interp-tsfc does with its own siblings. Revert this commit once the FIAT stack lands. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_013NoTXWyj2fVdJTHnMDFB4k --- .github/actions/install/action.yml | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/.github/actions/install/action.yml b/.github/actions/install/action.yml index e71d9f4f27..e6c3623350 100644 --- a/.github/actions/install/action.yml +++ b/.github/actions/install/action.yml @@ -165,17 +165,13 @@ runs: --extra-index-url https://download.pytorch.org/whl/cpu \ "./firedrake-repo[${{ inputs.deps }}]" - firedrake-clean - pip list - - - name: 'DROP BEFORE MERGE: build FIAT from the PR stack' - shell: bash - run: | - . venv/bin/activate - : # Head of the stack firedrakeproject/fiat#282 -> #284 -> #281 -> #286 + : # DROP BEFORE MERGE: the TSFC changes here need the GEM changes in the + : # FIAT stack firedrakeproject/fiat#282 -> #284 -> #281 -> #286, whose + : # head carries all four. This has to land before anything imports + : # Firedrake, firedrake-clean below included. pip install --no-deps --force-reinstall --ignore-installed \ git+https://github.com/firedrakeproject/fiat.git@pbrubeck/coffee-scalar-factor - : # The code generators changed, so discard kernels built against main + firedrake-clean pip list From 07a0c2aa4d00a91f92e1c82e619f3069c32721d8 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 29 Aug 2026 02:09:37 +0100 Subject: [PATCH 18/21] DROP BEFORE MERGE: install this branch's own FIAT sibling The stack head installed here, pbrubeck/coffee-scalar-factor, carries the GEM changes the TSFC half needs but not the FInAT half of this PR, so the job died importing Firedrake: ImportError: cannot import name 'CollapsedTensorProductPointSet' from 'finat.point_set' firedrakeproject/fiat#262 sits on top of that head and carries both. Install it instead. Revert this commit once the FIAT stack lands. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_013NoTXWyj2fVdJTHnMDFB4k --- .github/actions/install/action.yml | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/.github/actions/install/action.yml b/.github/actions/install/action.yml index e6c3623350..489e1c28a0 100644 --- a/.github/actions/install/action.yml +++ b/.github/actions/install/action.yml @@ -165,12 +165,12 @@ runs: --extra-index-url https://download.pytorch.org/whl/cpu \ "./firedrake-repo[${{ inputs.deps }}]" - : # DROP BEFORE MERGE: the TSFC changes here need the GEM changes in the - : # FIAT stack firedrakeproject/fiat#282 -> #284 -> #281 -> #286, whose - : # head carries all four. This has to land before anything imports - : # Firedrake, firedrake-clean below included. + : # DROP BEFORE MERGE: the changes here need the FInAT and GEM changes + : # in the FIAT stack firedrakeproject/fiat#282 -> #284 -> #281 -> + : # #286 -> #262, whose head carries all five. This has to land + : # before anything imports Firedrake, firedrake-clean below included. pip install --no-deps --force-reinstall --ignore-installed \ - git+https://github.com/firedrakeproject/fiat.git@pbrubeck/coffee-scalar-factor + git+https://github.com/firedrakeproject/fiat.git@pbrubeck/simplex-sum-factor firedrake-clean pip list From 4301517d8fc23f25bf45d3c62a3ebff864f14a14 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 29 Aug 2026 10:37:47 +0100 Subject: [PATCH 19/21] DROP BEFORE MERGE: install the FIAT stack in the docs job too The docs job installs Firedrake itself rather than going through .github/actions/install, so the stack this branch needs never reached it, and it failed the same way the test jobs did before 34c5f58: ImportError: cannot import name 'eliminate_deltas' from 'gem.optimise' raised by firedrake-clean at the end of its own install step. Give it its own copy of the install, in the same place relative to firedrake-clean. Revert this along with 34c5f58 once the FIAT stack lands. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_013NoTXWyj2fVdJTHnMDFB4k --- .github/workflows/core.yml | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/.github/workflows/core.yml b/.github/workflows/core.yml index 2e1ea2735f..2accda6dcd 100644 --- a/.github/workflows/core.yml +++ b/.github/workflows/core.yml @@ -567,6 +567,14 @@ jobs: pip install --verbose -r ./firedrake-repo/requirements-build.txt CC=mpicc CXX=mpicxx \ pip install --verbose --no-build-isolation './firedrake-repo[docs]' + + : # DROP BEFORE MERGE: this job installs Firedrake itself rather than + : # going through .github/actions/install, so it needs its own copy of + : # the FIAT stack, and for the same reason: firedrake-clean below + : # imports Firedrake, and so the tsfc that needs those GEM changes. + pip install --no-deps --force-reinstall --ignore-installed \ + git+https://github.com/firedrakeproject/fiat.git@pbrubeck/coffee-scalar-factor + firedrake-clean pip list From 67b70860c0c780901a15e8ec26e132db3de6a953 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 29 Aug 2026 15:53:54 +0100 Subject: [PATCH 20/21] Follow the renamed GEM delta and cost entry points FIAT renamed eliminate_deltas to cancel_nested_deltas, and moved the cost model to gem.cost. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_013NoTXWyj2fVdJTHnMDFB4k --- tsfc/spectral.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tsfc/spectral.py b/tsfc/spectral.py index bb9665eefa..628cbcff95 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -4,7 +4,8 @@ from gem.gem import Delta, Indexed, Sum, index_sum, one from gem.node import Memoizer, MemoizerArg -from gem.optimise import (eliminate_deltas, estimate_cost, +from gem.cost import estimate_cost +from gem.optimise import (cancel_nested_deltas, tabulate_indirect_contractions, filtered_replace_indices, has_linear_maps) from gem.optimise import delta_elimination as _delta_elimination @@ -38,7 +39,7 @@ def Integrals(expressions, quadrature_multiindex, argument_multiindices, paramet # Cancel the Deltas that select a basis transformation's columns, so that # monomial collection sees the resulting gather rather than the Delta. - expressions = [eliminate_deltas(e) for e in expressions] + expressions = [cancel_nested_deltas(e) for e in expressions] # Unroll max_extent = parameters["unroll_indexsum"] From c8b7043b4340a8d9aa097fdfe276f63cb1597fe2 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 29 Aug 2026 15:54:35 +0100 Subject: [PATCH 21/21] Import the GEM pipelines and lattice helpers from their own modules FIAT split gem.optimise. The pipelines that compose several passes are in gem.driver, and everything about a jagged lattice is in gem.jagged. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_013NoTXWyj2fVdJTHnMDFB4k --- tsfc/fem.py | 3 ++- tsfc/kernel_interface/common.py | 3 ++- tsfc/loopy.py | 5 +++-- 3 files changed, 7 insertions(+), 4 deletions(-) diff --git a/tsfc/fem.py b/tsfc/fem.py index 943089052e..d1beaec041 100644 --- a/tsfc/fem.py +++ b/tsfc/fem.py @@ -16,6 +16,7 @@ from finat.point_set import PointSet, PointSingleton from finat.quadrature import make_quadrature from finat.element_factory import as_fiat_cell, create_element +from gem.driver import contraction from gem.node import traversal from gem.optimise import constant_fold_zero, ffc_rounding from gem.unconcatenate import unconcatenate @@ -781,7 +782,7 @@ def take_singleton(xs): for var, expr in unconcatenate([(vec_beta, table_qi)], ctx.index_cache): indices = tuple(i for i in var.index_ordering() if i not in ctx.unsummed_coefficient_indices) value = gem.IndexSum(gem.Product(expr, var), indices) - summands.append(gem.optimise.contraction(value)) + summands.append(contraction(value)) optimised_value = gem.optimise.make_sum(summands) value_dict[alpha] = gem.ComponentTensor(optimised_value, zeta) diff --git a/tsfc/kernel_interface/common.py b/tsfc/kernel_interface/common.py index a807b3a161..0d95a8e075 100644 --- a/tsfc/kernel_interface/common.py +++ b/tsfc/kernel_interface/common.py @@ -18,7 +18,8 @@ from finat.duffy import DuffyElement from finat.quadrature import AbstractQuadratureRule from gem.node import traversal -from gem.optimise import constant_fold_zero, unflatten_returns +from gem.driver import unflatten_returns +from gem.optimise import constant_fold_zero from gem.optimise import remove_componenttensors as prune from numpy import asarray from tsfc import fem diff --git a/tsfc/loopy.py b/tsfc/loopy.py index c71372c018..abba465ff6 100644 --- a/tsfc/loopy.py +++ b/tsfc/loopy.py @@ -8,6 +8,7 @@ from collections import defaultdict, OrderedDict from gem import gem, impero as imp +from gem.jagged import compact_index_layout, simplex_lattice_ranks from gem.node import Memoizer import islpy as isl @@ -216,7 +217,7 @@ def lattice_rank(self, component): except KeyError: table = p.Variable(self.name_gen("lattice_rank")) self.lattice_ranks[shape] = ( - table, gem.simplex_lattice_ranks(component)) + table, simplex_lattice_ranks(component)) return p.Subscript(table, tuple(self.active_indices[index] for index in component)) @@ -281,7 +282,7 @@ def generate(impero_c, args, scalar_type, kernel_name="loopy_kernel", index_name if isinstance(temp, gem.Constant): shape, layout = temp.shape, None else: - shape, layout = gem.compact_index_layout( + shape, layout = compact_index_layout( tuple(ctx.indices[temp])) shape += temp.shape descriptors.append((temp, dtype, shape, layout))