Merge pull request #3904 from firedrakeproject/connorjward/add-tsfc

Absorb TSFC into Firedrake

.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,32 +84,40 @@ 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
           export CC=mpicc
           export HDF5_DIR="$PETSC_DIR/$PETSC_ARCH"
           export HDF5_MPI=ON
           python -m pip install --no-binary h5py -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,35 @@ 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")])
-
-projects = dict(
-    [("firedrake", "firedrakeproject"),
-     ("tsfc", "firedrakeproject"),
-     ("ufl", "firedrakeproject"),
-     ("FInAT", "FInAT"),
-     ("fiat", "firedrakeproject"),
-     ("petsc", "firedrakeproject"),
-     ("petsc4py", "firedrakeproject"),
-     ("loopy", "firedrakeproject"),
-     ("slepc", "firedrakeproject"),
-     ("slepc4py", "firedrakeproject")])
+    # removed components, left so old Firedrake versions are archivable
+    ("PyOP2", "Framework for performance-portable parallel computations on unstructured meshes"),
+    ("tsfc", "The Two Stage Form Compiler"),
+    ("FInAT", "a smarter library of finite elements"),
+])
+
+projects = dict([
+    ("firedrake", "firedrakeproject"),
+    ("ufl", "firedrakeproject"),
+    ("fiat", "firedrakeproject"),
+    ("petsc", "firedrakeproject"),
+    ("loopy", "firedrakeproject"),
+    ("slepc", "firedrakeproject"),
+    # removed components, left so old Firedrake versions are archivable
+    ("PyOP2", "OP2"),
+    ("tsfc", "firedrakeproject"),
+    ("FInAT", "FInAT"),
+])
 
 components = list(descriptions.keys())
 
-optional_components = ("slepc", "slepc4py", "petsc4py")
+optional_components = ("slepc", "PyOP2", "tsfc", "FInAT")
 
 parser = ArgumentParser(description="""Create Zenodo DOIs for specific versions of Firedrake components.
 

pyproject.toml

@@ -36,8 +36,6 @@ dependencies = [
   "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",
   "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/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

scripts/firedrake-install

@@ -1943,6 +1943,19 @@ else:
         if args.rebuild:
             pip_uninstall("petsc4py")
 
+    # Handle archived repositories
+    archived_repos = [("PyOP2", "firedrake", "PyOP2"),
+                      ("tsfc", "firedrake+fiat", "tsfc"),
+                      ("FInAT", "fiat", "FInAT")]
+    for src_repo, dest_repo, pip_pkg_name in archived_repos:
+        archived_path = os.path.join(firedrake_env, "src", src_repo)
+        if os.path.exists(archived_path):
+            log.warning("%s has been moved into %s, renaming to %s_old and pip uninstalling.\n"
+                        % (src_repo, dest_repo, src_repo))
+            pip_uninstall(pip_pkg_name)
+            new_path = os.path.join(firedrake_env, "src", "%s_old" % src_repo)
+            os.rename(archived_path, new_path)
+
     # If there is a petsc package to remove, then we're an old installation which will need to rebuild PETSc.
     update_petsc = pip_uninstall("petsc") or update_petsc
 

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:])