Add numba backend

.github/workflows/pythonapp.yml

@@ -122,8 +122,10 @@ jobs:
         uses: ilammy/msvc-dev-cmd@v1
 
       - name: Run FFCx demos
-        run: |
+        run: >
           pytest demo/test_demos.py
+          -W error
+          # -n auto
 
       - name: Build documentation
         run: |

demo/MassAction.py

@@ -5,7 +5,7 @@
 # SPDX-License-Identifier:    LGPL-3.0-or-later
 """Mass action demo."""
 
-import basix
+import basix.ufl
 import ufl
 
 P = 3

demo/test_demos.py

@@ -1,60 +1,87 @@
 """Test demos."""
 
 import os
+import subprocess
 import sys
+from pathlib import Path
 
 import pytest
 
-demo_dir = os.path.dirname(os.path.realpath(__file__))
+demo_dir = Path(__file__).parent
 
-ufl_files = []
-for file in os.listdir(demo_dir):
-    if file.endswith(".py") and not file == "test_demos.py":
-        ufl_files.append(file[:-3])
+ufl_files = [
+    f
+    for f in demo_dir.iterdir()
+    if f.suffix == ".py" and not f.stem.endswith("_numba") and f != Path(__file__)
+]
+
+skip_complex = ["BiharmonicHHJ", "BiharmonicRegge", "StabilisedStokes"]
+
+
+def skip_unsupported(test):
+    """Dcecorate test case to skip unsupported cases."""
+
+    def check_skip(file, scalar_type):
+        """Skip scalar_type file combinations not supported."""
+        if "complex" in scalar_type and file.stem in skip_complex:
+            pytest.skip(reason="Not implemented for complex types")
+        elif "Complex" in file.stem and scalar_type in ["float64", "float32"]:
+            pytest.skip(reason="Not implemented for real types")
+
+        return test(file, scalar_type)
+
+    return check_skip
 
 
 @pytest.mark.parametrize("file", ufl_files)
 @pytest.mark.parametrize("scalar_type", ["float64", "float32", "complex128", "complex64"])
-def test_demo(file, scalar_type):
+@skip_unsupported
+def test_C(file, scalar_type):
     """Test a demo."""
     if sys.platform.startswith("win32") and "complex" in scalar_type:
         # Skip complex demos on win32
         pytest.skip(reason="_Complex not supported on Windows")
 
-    if "complex" in scalar_type and file in [
-        "BiharmonicHHJ",
-        "BiharmonicRegge",
-        "StabilisedStokes",
-    ]:
-        # Skip demos that are not implemented for complex scalars
-        pytest.skip(reason="Not implemented for complex types")
-    elif "Complex" in file and scalar_type in ["float64", "float32"]:
-        # Skip demos that are only implemented for complex scalars
-        pytest.skip(reason="Not implemented for real types")
+    subprocess.run(["ffcx", "--scalar_type", scalar_type, file], cwd=demo_dir, check=True)
 
     if sys.platform.startswith("win32"):
-        opts = f"--scalar_type {scalar_type}"
         extra_flags = "/std:c17"
-        assert os.system(f"cd {demo_dir} && ffcx {opts} {file}.py") == 0
-        assert (
-            os.system(
-                f'cd {demo_dir} && cl.exe /I "../ffcx/codegeneration" {extra_flags} /c {file}.c'
+        for compiler in ["cl.exe", "clang-cl.exe"]:
+            subprocess.run(
+                [
+                    compiler,
+                    "/I",
+                    f"{demo_dir.parent / 'ffcx/codegeneration'}",
+                    *extra_flags.split(" "),
+                    "/c",
+                    file.with_suffix(".c"),
+                ],
+                cwd=demo_dir,
+                check=True,
             )
-        ) == 0
-        assert (
-            os.system(
-                f"cd {demo_dir} && "
-                f'clang-cl.exe /I "../ffcx/codegeneration" {extra_flags} /c {file}.c'
-            )
-        ) == 0
     else:
         cc = os.environ.get("CC", "cc")
-        opts = f"--scalar_type {scalar_type}"
         extra_flags = (
             "-std=c17 -Wunused-variable -Werror -fPIC -Wno-error=implicit-function-declaration"
         )
-        assert os.system(f"cd {demo_dir} && ffcx {opts} {file}.py") == 0
-        assert (
-            os.system(f"cd {demo_dir} && {cc} -I../ffcx/codegeneration {extra_flags} -c {file}.c")
-            == 0
+        subprocess.run(
+            [
+                cc,
+                f"-I{demo_dir.parent / 'ffcx/codegeneration'}",
+                *extra_flags.split(" "),
+                "-c",
+                file.with_suffix(".c"),
+            ],
+            cwd=demo_dir,
+            check=True,
         )
+
+
+@pytest.mark.parametrize("file", ufl_files)
+@pytest.mark.parametrize("scalar_type", ["float64", "float32", "complex128", "complex64"])
+@skip_unsupported
+def test_numba(file, scalar_type):
+    """Test numba generation."""
+    opts = f"--language numba --scalar_type {scalar_type}"
+    subprocess.run(["ffcx", *opts.split(" "), file], cwd=demo_dir, check=True)
+    subprocess.run(["python", file], cwd=demo_dir, check=True)

ffcx/codegeneration/C/__init__.py

@@ -1 +1,5 @@
 """Generation of C code."""
+
+from ffcx.codegeneration.C import expressions, file, form, integrals
+
+__all__ = ["expressions", "file", "form", "integrals", "suffixes"]

ffcx/codegeneration/C/expressions.py

@@ -8,12 +8,13 @@
 from __future__ import annotations
 
 import logging
+import string
 
 import numpy as np
 
 from ffcx.codegeneration.backend import FFCXBackend
 from ffcx.codegeneration.C import expressions_template
-from ffcx.codegeneration.C.c_implementation import CFormatter
+from ffcx.codegeneration.C.implementation import Formatter
 from ffcx.codegeneration.expression_generator import ExpressionGenerator
 from ffcx.codegeneration.utils import dtype_to_c_type, dtype_to_scalar_dtype
 from ffcx.ir.representation import ExpressionIR
@@ -43,8 +44,8 @@ def generator(ir: ExpressionIR, options):
     d["factory_name"] = factory_name
     parts = eg.generate()
 
-    CF = CFormatter(options["scalar_type"])
-    d["tabulate_expression"] = CF.c_format(parts)
+    CF = Formatter(options["scalar_type"])
+    d["tabulate_expression"] = CF.format(parts)
 
     if len(ir.original_coefficient_positions) > 0:
         d["original_coefficient_positions"] = f"original_coefficient_positions_{factory_name}"
@@ -106,9 +107,10 @@ def generator(ir: ExpressionIR, options):
     d["coordinate_element_hash"] = f"UINT64_C({ir.expression.coordinate_element_hash})"
 
     # Check that no keys are redundant or have been missed
-    from string import Formatter
 
-    fields = [fname for _, fname, _, _ in Formatter().parse(expressions_template.factory) if fname]
+    fields = [
+        fname for _, fname, _, _ in string.Formatter().parse(expressions_template.factory) if fname
+    ]
     assert set(fields) == set(d.keys()), "Mismatch between keys in template and in formatting dict"
 
     # Format implementation code

ffcx/codegeneration/C/c_implementation.py → ffcx/codegeneration/C/implementation.py

@@ -142,7 +142,7 @@
 }
 
 
-class CFormatter:
+class Formatter:
     """C formatter."""
 
     scalar_type: np.dtype
@@ -186,7 +186,7 @@ def _build_initializer_lists(self, values):
 
     def format_statement_list(self, slist) -> str:
         """Format a statement list."""
-        return "".join(self.c_format(s) for s in slist.statements)
+        return "".join(self.format(s) for s in slist.statements)
 
     def format_section(self, section) -> str:
         """Format a section."""
@@ -197,12 +197,12 @@ def format_section(self, section) -> str:
             f"// Inputs: {', '.join(w.name for w in section.input)}\n"
             f"// Outputs: {', '.join(w.name for w in section.output)}\n"
         )
-        declarations = "".join(self.c_format(s) for s in section.declarations)
+        declarations = "".join(self.format(s) for s in section.declarations)
 
         body = ""
         if len(section.statements) > 0:
             declarations += "{\n  "
-            body = "".join(self.c_format(s) for s in section.statements)
+            body = "".join(self.format(s) for s in section.statements)
             body = body.replace("\n", "\n  ")
             body = body[:-2] + "}\n"
 
@@ -218,7 +218,7 @@ def format_array_decl(self, arr) -> str:
         dtype = arr.symbol.dtype
         typename = self._dtype_to_name(dtype)
 
-        symbol = self.c_format(arr.symbol)
+        symbol = self.format(arr.symbol)
         dims = "".join([f"[{i}]" for i in arr.sizes])
         if arr.values is None:
             assert arr.const is False
@@ -230,21 +230,21 @@ def format_array_decl(self, arr) -> str:
 
     def format_array_access(self, arr) -> str:
         """Format an array access."""
-        name = self.c_format(arr.array)
-        indices = f"[{']['.join(self.c_format(i) for i in arr.indices)}]"
+        name = self.format(arr.array)
+        indices = f"[{']['.join(self.format(i) for i in arr.indices)}]"
         return f"{name}{indices}"
 
     def format_variable_decl(self, v) -> str:
         """Format a variable declaration."""
-        val = self.c_format(v.value)
-        symbol = self.c_format(v.symbol)
+        val = self.format(v.value)
+        symbol = self.format(v.symbol)
         typename = self._dtype_to_name(v.symbol.dtype)
         return f"{typename} {symbol} = {val};\n"
 
     def format_nary_op(self, oper) -> str:
         """Format an n-ary operation."""
         # Format children
-        args = [self.c_format(arg) for arg in oper.args]
+        args = [self.format(arg) for arg in oper.args]
 
         # Apply parentheses
         for i in range(len(args)):
@@ -257,8 +257,8 @@ def format_nary_op(self, oper) -> str:
     def format_binary_op(self, oper) -> str:
         """Format a binary operation."""
         # Format children
-        lhs = self.c_format(oper.lhs)
-        rhs = self.c_format(oper.rhs)
+        lhs = self.format(oper.lhs)
+        rhs = self.format(oper.rhs)
 
         # Apply parentheses
         if oper.lhs.precedence >= oper.precedence:
@@ -271,7 +271,7 @@ def format_binary_op(self, oper) -> str:
 
     def format_unary_op(self, oper) -> str:
         """Format a unary operation."""
-        arg = self.c_format(oper.arg)
+        arg = self.format(oper.arg)
         if oper.arg.precedence >= oper.precedence:
             return f"{oper.op}({arg})"
         return f"{oper.op}{arg}"
@@ -287,12 +287,12 @@ def format_literal_int(self, val) -> str:
 
     def format_for_range(self, r) -> str:
         """Format a for loop over a range."""
-        begin = self.c_format(r.begin)
-        end = self.c_format(r.end)
-        index = self.c_format(r.index)
+        begin = self.format(r.begin)
+        end = self.format(r.end)
+        index = self.format(r.index)
         output = f"for (int {index} = {begin}; {index} < {end}; ++{index})\n"
         output += "{\n"
-        body = self.c_format(r.body)
+        body = self.format(r.body)
         for line in body.split("\n"):
             if len(line) > 0:
                 output += f"  {line}\n"
@@ -301,20 +301,20 @@ def format_for_range(self, r) -> str:
 
     def format_statement(self, s) -> str:
         """Format a statement."""
-        return self.c_format(s.expr)
+        return self.format(s.expr)
 
     def format_assign(self, expr) -> str:
         """Format an assignment."""
-        rhs = self.c_format(expr.rhs)
-        lhs = self.c_format(expr.lhs)
+        rhs = self.format(expr.rhs)
+        lhs = self.format(expr.lhs)
         return f"{lhs} {expr.op} {rhs};\n"
 
     def format_conditional(self, s) -> str:
         """Format a conditional."""
         # Format children
-        c = self.c_format(s.condition)
-        t = self.c_format(s.true)
-        f = self.c_format(s.false)
+        c = self.format(s.condition)
+        t = self.format(s.true)
+        f = self.format(s.false)
 
         # Apply parentheses
         if s.condition.precedence >= s.precedence:
@@ -333,7 +333,7 @@ def format_symbol(self, s) -> str:
 
     def format_multi_index(self, mi) -> str:
         """Format a multi-index."""
-        return self.c_format(mi.global_index)
+        return self.format(mi.global_index)
 
     def format_math_function(self, c) -> str:
         """Format a mathematical function."""
@@ -349,10 +349,10 @@ def format_math_function(self, c) -> str:
 
         # Get a function from the table, if available, else just use bare name
         func = dtype_math_table.get(c.function, c.function)
-        args = ", ".join(self.c_format(arg) for arg in c.args)
+        args = ", ".join(self.format(arg) for arg in c.args)
         return f"{func}({args})"
 
-    c_impl = {
+    impl = {
         "Section": format_section,
         "StatementList": format_statement_list,
         "Comment": format_comment,
@@ -387,10 +387,10 @@ def format_math_function(self, c) -> str:
         "LT": format_binary_op,
     }
 
-    def c_format(self, s) -> str:
+    def format(self, s) -> str:
         """Format as C."""
         name = s.__class__.__name__
         try:
-            return self.c_impl[name](self, s)
+            return self.impl[name](self, s)
         except KeyError:
             raise RuntimeError("Unknown statement: ", name)