DO NOT SQUASH Absorb TSFC

.github/workflows/pyop2.yml

@@ -1,14 +1,10 @@
-name: PyOP2
+name: Test PyOP2 and TSFC
 
-# Trigger the workflow on push or pull request,
-# but only for the master branch
 on:
   push:
     branches:
       - master
   pull_request:
-    branches:
-      - master
 
 jobs:
   test:
@@ -88,29 +84,37 @@ jobs:
           make
           make install
 
-      - name: Checkout PyOP2
+      - name: Checkout Firedrake
         uses: actions/checkout@v4
         with:
-          path: PyOP2
+          path: firedrake
 
       - name: Install PyOP2 dependencies
         shell: bash
-        working-directory: PyOP2
+        working-directory: firedrake
         run: |
           source ../venv/bin/activate
           python -m pip install -U pip
           python -m pip install -U pytest-timeout
 
       - name: Install PyOP2
         shell: bash
-        working-directory: PyOP2
+        working-directory: firedrake
         run: |
           source ../venv/bin/activate
           python -m pip install -v ".[test]"
 
-      - name: Run tests
+      - name: Run TSFC tests
+        shell: bash
+        working-directory: firedrake
+        run: |
+          source ../venv/bin/activate
+          pytest --tb=native --timeout=480 --timeout-method=thread -o faulthandler_timeout=540 -v tests/tsfc
+        timeout-minutes: 10
+
+      - name: Run PyOP2 tests
         shell: bash
-        working-directory: PyOP2
+        working-directory: firedrake
         run: |
           source ../venv/bin/activate
           # Running parallel test cases separately works around a bug in pytest-mpi

Makefile

@@ -24,6 +24,8 @@ lint:
 	@python -m flake8 $(FLAKE8_FORMAT) pyop2
 	@echo "    Linting PyOP2 scripts"
 	@python -m flake8 $(FLAKE8_FORMAT) pyop2/scripts --filename=*
+	@echo "    Linting TSFC"
+	@python -m flake8 $(FLAKE8_FORMAT) tsfc
 
 actionlint:
 	@echo "    Pull latest actionlint image"

firedrake/scripts/firedrake-zenodo

@@ -16,35 +16,25 @@ import datetime
 
 # Change this to https://sandbox.zenodo.org/api for testing
 ZENODO_URL = "https://zenodo.org/api"
-# TODO: Remove "petsc4py" once all users have switched to "petsc/src/binding/petsc4py".
-# And the same for slepc4py.
 descriptions = OrderedDict([
     ("firedrake", "an automated finite element system"),
-    ("tsfc", "The Two Stage Form Compiler"),
     ("ufl", "The Unified Form Language"),
-    ("FInAT", "a smarter library of finite elements"),
     ("fiat", "The Finite Element Automated Tabulator"),
     ("petsc", "Portable, Extensible Toolkit for Scientific Computation"),
-    ("petsc4py", "The Python interface to PETSc"),
     ("loopy", "Transformation-Based Generation of High-Performance CPU/GPU Code"),
-    ("slepc", "Scalable Library for Eigenvalue Problem Computations"),
-    ("slepc4py", "The Python interface to SLEPc")])
+    ("slepc", "Scalable Library for Eigenvalue Problem Computations")])
 
 projects = dict(
     [("firedrake", "firedrakeproject"),
-     ("tsfc", "firedrakeproject"),
      ("ufl", "firedrakeproject"),
-     ("FInAT", "FInAT"),
      ("fiat", "firedrakeproject"),
      ("petsc", "firedrakeproject"),
-     ("petsc4py", "firedrakeproject"),
      ("loopy", "firedrakeproject"),
-     ("slepc", "firedrakeproject"),
-     ("slepc4py", "firedrakeproject")])
+     ("slepc", "firedrakeproject")])
 
 components = list(descriptions.keys())
 
-optional_components = ("slepc", "slepc4py", "petsc4py")
+optional_components = ("slepc",)
 
 parser = ArgumentParser(description="""Create Zenodo DOIs for specific versions of Firedrake components.
 

pyproject.toml

@@ -30,9 +30,7 @@ dependencies = [
   "scipy",
   "sympy",
   "fenics-ufl @ git+https://github.com/firedrakeproject/ufl.git",
-  "fenics-fiat @ git+https://github.com/firedrakeproject/fiat.git",
-  "finat @ git+https://github.com/FInAT/FInAT.git",
-  "tsfc @ git+https://github.com/firedrakeproject/tsfc.git",
+  "fenics-fiat @ git+https://github.com/firedrakeproject/fiat.git@connorjward/add-finat-gem",
   "pyadjoint-ad @ git+https://github.com/dolfin-adjoint/pyadjoint.git",
   "loopy @ git+https://github.com/firedrakeproject/loopy.git@main",
 ]

requirements-git.txt

@@ -1,7 +1,5 @@
 git+https://github.com/firedrakeproject/ufl.git#egg=fenics-ufl
-git+https://github.com/firedrakeproject/fiat.git#egg=fenics-fiat
-git+https://github.com/FInAT/FInAT.git#egg=finat
-git+https://github.com/firedrakeproject/tsfc.git#egg=tsfc
+git+https://github.com/firedrakeproject/fiat.git@connorjward/add-finat-gem#egg=fenics-fiat
 git+https://github.com/dolfin-adjoint/pyadjoint.git#egg=pyadjoint-ad
 git+https://github.com/firedrakeproject/loopy.git@main#egg=loopy
 git+https://github.com/firedrakeproject/pytest-mpi.git@main#egg=pytest-mpi

tests/tsfc/test_codegen.py

@@ -0,0 +1,30 @@
+import pytest
+
+from gem import impero_utils
+from gem.gem import Index, Indexed, IndexSum, Product, Variable
+
+
+def test_loop_fusion():
+    i = Index()
+    j = Index()
+    Ri = Indexed(Variable('R', (6,)), (i,))
+
+    def make_expression(i, j):
+        A = Variable('A', (6,))
+        s = IndexSum(Indexed(A, (j,)), (j,))
+        return Product(Indexed(A, (i,)), s)
+
+    e1 = make_expression(i, j)
+    e2 = make_expression(i, i)
+
+    def gencode(expr):
+        impero_c = impero_utils.compile_gem([(Ri, expr)], (i, j))
+        return impero_c.tree
+
+    assert len(gencode(e1).children) == len(gencode(e2).children)
+
+
+if __name__ == "__main__":
+    import os
+    import sys
+    pytest.main(args=[os.path.abspath(__file__)] + sys.argv[1:])

tests/tsfc/test_coffee_optimise.py

@@ -0,0 +1,78 @@
+import pytest
+
+from gem.gem import Index, Indexed, Product, Variable, Division, Literal, Sum
+from gem.optimise import replace_division
+from tsfc.coffee_mode import optimise_expressions
+
+
+def test_replace_div():
+    i = Index()
+    A = Variable('A', ())
+    B = Variable('B', (6,))
+    Bi = Indexed(B, (i,))
+    d = Division(Bi, A)
+    result, = replace_division([d])
+    expected = Product(Bi, Division(Literal(1.0), A))
+
+    assert result == expected
+
+
+def test_loop_optimise():
+    I = 20
+    J = K = 10
+    i = Index('i', I)
+    j = Index('j', J)
+    k = Index('k', K)
+
+    A1 = Variable('a1', (I,))
+    A2 = Variable('a2', (I,))
+    A3 = Variable('a3', (I,))
+    A1i = Indexed(A1, (i,))
+    A2i = Indexed(A2, (i,))
+    A3i = Indexed(A3, (i,))
+
+    B = Variable('b', (J,))
+    C = Variable('c', (J,))
+    Bj = Indexed(B, (j,))
+    Cj = Indexed(C, (j,))
+
+    E = Variable('e', (K,))
+    F = Variable('f', (K,))
+    G = Variable('g', (K,))
+    Ek = Indexed(E, (k,))
+    Fk = Indexed(F, (k,))
+    Gk = Indexed(G, (k,))
+
+    Z = Variable('z', ())
+
+    # Bj*Ek + Bj*Fk => (Ek + Fk)*Bj
+    expr = Sum(Product(Bj, Ek), Product(Bj, Fk))
+    result, = optimise_expressions([expr], (j, k))
+    expected = Product(Sum(Ek, Fk), Bj)
+    assert result == expected
+
+    # Bj*Ek + Bj*Fk + Bj*Gk + Cj*Ek + Cj*Fk =>
+    # (Ek + Fk + Gk)*Bj + (Ek+Fk)*Cj
+    expr = Sum(Sum(Sum(Sum(Product(Bj, Ek), Product(Bj, Fk)), Product(Bj, Gk)),
+                   Product(Cj, Ek)), Product(Cj, Fk))
+    result, = optimise_expressions([expr], (j, k))
+    expected = Sum(Product(Sum(Sum(Ek, Fk), Gk), Bj), Product(Sum(Ek, Fk), Cj))
+    assert result == expected
+
+    # Z*A1i*Bj*Ek + Z*A2i*Bj*Ek + A3i*Bj*Ek + Z*A1i*Bj*Fk =>
+    # Bj*(Ek*(Z*A1i + Z*A2i) + A3i) + Z*A1i*Fk)
+
+    expr = Sum(Sum(Sum(Product(Z, Product(A1i, Product(Bj, Ek))),
+                       Product(Z, Product(A2i, Product(Bj, Ek)))),
+                   Product(A3i, Product(Bj, Ek))),
+               Product(Z, Product(A1i, Product(Bj, Fk))))
+    result, = optimise_expressions([expr], (j, k))
+    expected = Product(Sum(Product(Ek, Sum(Sum(Product(Z, A1i), Product(Z, A2i)), A3i)),
+                           Product(Fk, Product(Z, A1i))), Bj)
+    assert result == expected
+
+
+if __name__ == "__main__":
+    import os
+    import sys
+    pytest.main(args=[os.path.abspath(__file__)] + sys.argv[1:])

tests/tsfc/test_create_fiat_element.py

@@ -0,0 +1,150 @@
+import pytest
+
+import FIAT
+from FIAT.discontinuous_lagrange import DiscontinuousLagrange as FIAT_DiscontinuousLagrange
+
+import ufl
+import finat.ufl
+from tsfc.finatinterface import create_element as _create_element
+
+
+supported_elements = {
+    # These all map directly to FIAT elements
+    "Brezzi-Douglas-Marini": FIAT.BrezziDouglasMarini,
+    "Brezzi-Douglas-Fortin-Marini": FIAT.BrezziDouglasFortinMarini,
+    "Lagrange": FIAT.Lagrange,
+    "Nedelec 1st kind H(curl)": FIAT.Nedelec,
+    "Nedelec 2nd kind H(curl)": FIAT.NedelecSecondKind,
+    "Raviart-Thomas": FIAT.RaviartThomas,
+    "Regge": FIAT.Regge,
+}
+"""A :class:`.dict` mapping UFL element family names to their
+FIAT-equivalent constructors."""
+
+
+def create_element(ufl_element):
+    """Create a FIAT element given a UFL element."""
+    finat_element = _create_element(ufl_element)
+    return finat_element.fiat_equivalent
+
+
+@pytest.fixture(params=["BDM",
+                        "BDFM",
+                        "Lagrange",
+                        "N1curl",
+                        "N2curl",
+                        "RT",
+                        "Regge"])
+def triangle_names(request):
+    return request.param
+
+
+@pytest.fixture
+def ufl_element(triangle_names):
+    return finat.ufl.FiniteElement(triangle_names, ufl.triangle, 2)
+
+
+def test_triangle_basic(ufl_element):
+    element = create_element(ufl_element)
+    assert isinstance(element, supported_elements[ufl_element.family()])
+
+
+@pytest.fixture(params=["CG", "DG", "DG L2"], scope="module")
+def tensor_name(request):
+    return request.param
+
+
+@pytest.fixture(params=[ufl.interval, ufl.triangle,
+                        ufl.quadrilateral],
+                ids=lambda x: x.cellname(),
+                scope="module")
+def ufl_A(request, tensor_name):
+    return finat.ufl.FiniteElement(tensor_name, request.param, 1)
+
+
+@pytest.fixture
+def ufl_B(tensor_name):
+    return finat.ufl.FiniteElement(tensor_name, ufl.interval, 1)
+
+
+def test_tensor_prod_simple(ufl_A, ufl_B):
+    tensor_ufl = finat.ufl.TensorProductElement(ufl_A, ufl_B)
+
+    tensor = create_element(tensor_ufl)
+    A = create_element(ufl_A)
+    B = create_element(ufl_B)
+
+    assert isinstance(tensor, FIAT.TensorProductElement)
+
+    assert tensor.A is A
+    assert tensor.B is B
+
+
+@pytest.mark.parametrize(('family', 'expected_cls'),
+                         [('P', FIAT.GaussLobattoLegendre),
+                          ('DP', FIAT.GaussLegendre),
+                          ('DP L2', FIAT.GaussLegendre)])
+def test_interval_variant_default(family, expected_cls):
+    ufl_element = finat.ufl.FiniteElement(family, ufl.interval, 3)
+    assert isinstance(create_element(ufl_element), expected_cls)
+
+
+@pytest.mark.parametrize(('family', 'variant', 'expected_cls'),
+                         [('P', 'equispaced', FIAT.Lagrange),
+                          ('P', 'spectral', FIAT.GaussLobattoLegendre),
+                          ('DP', 'equispaced', FIAT_DiscontinuousLagrange),
+                          ('DP', 'spectral', FIAT.GaussLegendre),
+                          ('DP L2', 'equispaced', FIAT_DiscontinuousLagrange),
+                          ('DP L2', 'spectral', FIAT.GaussLegendre)])
+def test_interval_variant(family, variant, expected_cls):
+    ufl_element = finat.ufl.FiniteElement(family, ufl.interval, 3, variant=variant)
+    assert isinstance(create_element(ufl_element), expected_cls)
+
+
+def test_triangle_variant_spectral():
+    ufl_element = finat.ufl.FiniteElement('DP', ufl.triangle, 2, variant='spectral')
+    create_element(ufl_element)
+
+
+def test_triangle_variant_spectral_l2():
+    ufl_element = finat.ufl.FiniteElement('DP L2', ufl.triangle, 2, variant='spectral')
+    create_element(ufl_element)
+
+
+def test_quadrilateral_variant_spectral_q():
+    element = create_element(finat.ufl.FiniteElement('Q', ufl.quadrilateral, 3, variant='spectral'))
+    assert isinstance(element.element.A, FIAT.GaussLobattoLegendre)
+    assert isinstance(element.element.B, FIAT.GaussLobattoLegendre)
+
+
+def test_quadrilateral_variant_spectral_dq():
+    element = create_element(finat.ufl.FiniteElement('DQ', ufl.quadrilateral, 1, variant='spectral'))
+    assert isinstance(element.element.A, FIAT.GaussLegendre)
+    assert isinstance(element.element.B, FIAT.GaussLegendre)
+
+
+def test_quadrilateral_variant_spectral_dq_l2():
+    element = create_element(finat.ufl.FiniteElement('DQ L2', ufl.quadrilateral, 1, variant='spectral'))
+    assert isinstance(element.element.A, FIAT.GaussLegendre)
+    assert isinstance(element.element.B, FIAT.GaussLegendre)
+
+
+def test_quadrilateral_variant_spectral_rtcf():
+    element = create_element(finat.ufl.FiniteElement('RTCF', ufl.quadrilateral, 2, variant='spectral'))
+    assert isinstance(element.element._elements[0].A, FIAT.GaussLobattoLegendre)
+    assert isinstance(element.element._elements[0].B, FIAT.GaussLegendre)
+    assert isinstance(element.element._elements[1].A, FIAT.GaussLegendre)
+    assert isinstance(element.element._elements[1].B, FIAT.GaussLobattoLegendre)
+
+
+def test_cache_hit(ufl_element):
+    A = create_element(ufl_element)
+    B = create_element(ufl_element)
+
+    assert A is B
+
+
+if __name__ == "__main__":
+    import os
+    import sys
+    pytest.main(args=[os.path.abspath(__file__)] + sys.argv[1:])