Preserve finite element pullbacks when collecting monomials

.github/actions/install/action.yml

@@ -165,6 +165,13 @@ runs:
           --extra-index-url https://download.pytorch.org/whl/cpu \
           "./firedrake-repo[${{ inputs.deps }}]"
 
+        : # DROP BEFORE MERGE: the 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
+
         firedrake-clean
         pip list
 

.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
 

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)

tests/tsfc/test_sum_factorisation.py

@@ -3,10 +3,12 @@
 
 from ufl import (Mesh, FunctionSpace, TestFunction, TrialFunction,
                  TensorProductCell, dx, action, interval, triangle,
-                 quadrilateral, hexahedron, curl, dot, div, grad)
+                 quadrilateral, hexahedron, tetrahedron, curl, dot, div,
+                 grad, inner)
 from finat.ufl import (FiniteElement, VectorElement, EnrichedElement,
                        TensorProductElement, HCurlElement, HDivElement)
 
+import tsfc.spectral
 from tsfc import compile_form
 
 
@@ -79,6 +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),
@@ -203,6 +210,66 @@ def test_equivalent_cells(cell, equivalent_cell, degree):
     assert count_flops(action(a)) == count_flops(action(b))
 
 
+@pytest.fixture
+def expanded(monkeypatch):
+    """Force the expanded representation, for comparison."""
+    def force(monkeypatch=monkeypatch):
+        collect_monomials = tsfc.spectral.collect_monomials
+        monkeypatch.setattr(
+            tsfc.spectral, "collect_monomials",
+            lambda expressions, classifier, _: collect_monomials(
+                expressions, classifier))
+    return force
+
+
+def piola_helmholtz(cell, degree):
+    m = Mesh(VectorElement('CG', cell, 1))
+    V = FunctionSpace(m, FiniteElement('RT', cell, degree))
+    u = TrialFunction(V)
+    v = TestFunction(V)
+    return (inner(u, v) + inner(div(u), div(v)))*dx
+
+
+@pytest.mark.parametrize('cell', [triangle, tetrahedron],
+                         ids=lambda cell: cell.cellname)
+@pytest.mark.parametrize('degree', [1, 2, 3])
+def test_piola_map_is_preserved(cell, degree, expanded):
+    # Test and trial apply the same Piola map, so preserving it evaluates
+    # the physical basis once instead of pushing the geometry through both
+    # argument axes.
+    form = piola_helmholtz(cell, degree)
+    selected = count_flops(form)
+    expanded()
+    assert selected < count_flops(form)
+
+
+@pytest.mark.parametrize('cell', [triangle, tetrahedron],
+                         ids=lambda cell: cell.cellname)
+@pytest.mark.parametrize('degree', [1, 3])
+def test_preserving_a_map_is_never_worse(cell, degree, expanded):
+    # Expanding a map exposes scalar factorisation of its entries, which at
+    # some degrees beats sharing it.  Selection costs both, so neither
+    # representation may regress the other.
+    form = helmholtz(cell, degree)
+    selected = count_flops(form)
+    expanded()
+    assert selected <= count_flops(form)
+
+
+@pytest.mark.parametrize('cell', [triangle, tetrahedron],
+                         ids=lambda cell: cell.cellname)
+@pytest.mark.parametrize('degree', [1, 3])
+def test_shared_map_is_tabulated_in_one_loop(cell, degree, expanded):
+    # A map that both argument axes share is tabulated once, so it must be
+    # tabulated in one loop.  An index per axis fissions the loop nest that
+    # the expanded representation keeps whole, which costs more than the
+    # flops it saves.
+    form = piola_helmholtz(cell, degree)
+    selected = count_loops(form)
+    expanded()
+    assert selected <= count_loops(form)
+
+
 if __name__ == "__main__":
     import os
     import sys

tsfc/coffee_mode.py

@@ -2,7 +2,8 @@
 
 from gem.node import traversal, Memoizer
 from gem.gem import Failure, Sum, index_sum
-from gem.optimise import replace_division, unroll_indexsum
+from gem.optimise import (tabulate_indirect_contractions, replace_division,
+                          unroll_indexsum)
 from gem.refactorise import collect_monomials
 from gem.unconcatenate import unconcatenate
 from gem.coffee import optimise_monomial_sum
@@ -78,4 +79,5 @@ def optimise_expressions(expressions, argument_indices):
     classifier = partial(spectral.classify, set(argument_indices),
                          delta_inside=Memoizer(spectral._delta_inside))
     monomial_sums = collect_monomials(expressions, classifier)
-    return [optimise_monomial_sum(ms, argument_indices) for ms in monomial_sums]
+    return [tabulate_indirect_contractions(
+        optimise_monomial_sum(ms, argument_indices)) for ms in monomial_sums]

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:
@@ -547,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