diff --git a/.github/workflows/pythonapp.yml b/.github/workflows/pythonapp.yml index 216c0cb76..0b728cf12 100644 --- a/.github/workflows/pythonapp.yml +++ b/.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: | diff --git a/demo/MassAction.py b/demo/MassAction.py index 59f42ecd3..17ca2e868 100644 --- a/demo/MassAction.py +++ b/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 diff --git a/demo/test_demos.py b/demo/test_demos.py index bebfec2a2..c376b09a0 100644 --- a/demo/test_demos.py +++ b/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): + """Decorate 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) diff --git a/ffcx/codegeneration/C/__init__.py b/ffcx/codegeneration/C/__init__.py index e4d6c1f41..0cc7d4a8b 100644 --- a/ffcx/codegeneration/C/__init__.py +++ b/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"] diff --git a/ffcx/codegeneration/C/expressions.py b/ffcx/codegeneration/C/expressions.py index 177e8a0be..7ad470ccf 100644 --- a/ffcx/codegeneration/C/expressions.py +++ b/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 diff --git a/ffcx/codegeneration/C/c_implementation.py b/ffcx/codegeneration/C/implementation.py similarity index 89% rename from ffcx/codegeneration/C/c_implementation.py rename to ffcx/codegeneration/C/implementation.py index 1d6eafa23..6bdde8eec 100644 --- a/ffcx/codegeneration/C/c_implementation.py +++ b/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) diff --git a/ffcx/codegeneration/C/integrals.py b/ffcx/codegeneration/C/integrals.py index 9fb64ccc7..279121e29 100644 --- a/ffcx/codegeneration/C/integrals.py +++ b/ffcx/codegeneration/C/integrals.py @@ -10,10 +10,11 @@ import basix import numpy as np +from numpy import typing as npt from ffcx.codegeneration.backend import FFCXBackend from ffcx.codegeneration.C import integrals_template as ufcx_integrals -from ffcx.codegeneration.C.c_implementation import CFormatter +from ffcx.codegeneration.C.implementation import Formatter from ffcx.codegeneration.integral_generator import IntegralGenerator from ffcx.codegeneration.utils import dtype_to_c_type, dtype_to_scalar_dtype from ffcx.ir.representation import IntegralIR @@ -21,8 +22,20 @@ logger = logging.getLogger("ffcx") -def generator(ir: IntegralIR, domain: basix.CellType, options): - """Generate C code for an integral.""" +def generator( + ir: IntegralIR, domain: basix.CellType, options: dict[str, int | float | npt.DTypeLike] +) -> tuple[str, str]: + """Generate C code for an integral. + + Args: + ir: IR of the integral + domain: basix cell type + options: dict of kernel generation options + + Returns: + Tuple of declaration (header) and implementation (source) strings. + + """ logger.info("Generating code for integral:") logger.info(f"--- type: {ir.expression.integral_type}") logger.info(f"--- name: {ir.expression.name}") @@ -32,7 +45,7 @@ def generator(ir: IntegralIR, domain: basix.CellType, options): # Format declaration declaration = ufcx_integrals.declaration.format(factory_name=factory_name) - # Create FFCx C backend + # Create FFCx backend backend = FFCXBackend(ir, options) # Configure kernel generator @@ -42,8 +55,8 @@ def generator(ir: IntegralIR, domain: basix.CellType, options): parts = ig.generate(domain) # Format code as string - CF = CFormatter(options["scalar_type"]) - body = CF.c_format(parts) + CF = Formatter(options["scalar_type"]) # type: ignore + body = CF.format(parts) # Generate generic FFCx code snippets and add specific parts code = {} @@ -69,7 +82,7 @@ def generator(ir: IntegralIR, domain: basix.CellType, options): else: code["tabulate_tensor_complex64"] = ".tabulate_tensor_complex64 = NULL," code["tabulate_tensor_complex128"] = ".tabulate_tensor_complex128 = NULL," - np_scalar_type = np.dtype(options["scalar_type"]).name + np_scalar_type = np.dtype(options["scalar_type"]).name # type: ignore code[f"tabulate_tensor_{np_scalar_type}"] = ( f".tabulate_tensor_{np_scalar_type} = tabulate_tensor_{factory_name}," ) @@ -81,8 +94,8 @@ def generator(ir: IntegralIR, domain: basix.CellType, options): enabled_coefficients_init=code["enabled_coefficients_init"], tabulate_tensor=code["tabulate_tensor"], needs_facet_permutations="true" if ir.expression.needs_facet_permutations else "false", - scalar_type=dtype_to_c_type(options["scalar_type"]), - geom_type=dtype_to_c_type(dtype_to_scalar_dtype(options["scalar_type"])), + scalar_type=dtype_to_c_type(options["scalar_type"]), # type: ignore + geom_type=dtype_to_c_type(dtype_to_scalar_dtype(options["scalar_type"])), # type: ignore coordinate_element_hash=f"UINT64_C({ir.expression.coordinate_element_hash})", tabulate_tensor_float32=code["tabulate_tensor_float32"], tabulate_tensor_float64=code["tabulate_tensor_float64"], @@ -91,4 +104,6 @@ def generator(ir: IntegralIR, domain: basix.CellType, options): domain=int(domain), ) + # TODO: Check that no keys are redundant or have been missed (ref. numba/integrals.py) + return declaration, implementation diff --git a/ffcx/codegeneration/codegeneration.py b/ffcx/codegeneration/codegeneration.py index bc1ecfe2a..053fea70c 100644 --- a/ffcx/codegeneration/codegeneration.py +++ b/ffcx/codegeneration/codegeneration.py @@ -14,13 +14,10 @@ import logging import typing +from importlib import import_module import numpy.typing as npt -from ffcx.codegeneration.C.expressions import generator as expression_generator -from ffcx.codegeneration.C.file import generator as file_generator -from ffcx.codegeneration.C.form import generator as form_generator -from ffcx.codegeneration.C.integrals import generator as integral_generator from ffcx.ir.representation import DataIR logger = logging.getLogger("ffcx") @@ -45,6 +42,14 @@ def generate_code(ir: DataIR, options: dict[str, int | float | npt.DTypeLike]) - logger.info("Compiler stage 3: Generating code") logger.info(79 * "*") + lang = options.get("language", "C") + mod = import_module(f"ffcx.codegeneration.{lang}") + + integral_generator = mod.integrals.generator + form_generator = mod.form.generator + expression_generator = mod.expressions.generator + file_generator = mod.file.generator + code_integrals = [ integral_generator(integral_ir, domain, options) for integral_ir in ir.integrals diff --git a/ffcx/codegeneration/lnodes.py b/ffcx/codegeneration/lnodes.py index e846b0579..debd4441e 100644 --- a/ffcx/codegeneration/lnodes.py +++ b/ffcx/codegeneration/lnodes.py @@ -268,18 +268,24 @@ def __rdiv__(self, other): class LExprOperator(LExpr): """Base class for all expression operators.""" + precedence: int + sideeffect = False class LExprTerminal(LExpr): """Base class for all expression terminals.""" + precedence: int + sideeffect = False class LiteralFloat(LExprTerminal): """A floating point literal value.""" + precedence: int + precedence = PRECEDENCE.LITERAL def __init__(self, value): @@ -446,6 +452,8 @@ def __eq__(self, other): class BinOp(LExprOperator): """A binary operator.""" + op: str + def __init__(self, lhs, rhs): """Initialise.""" self.lhs = as_lexpr(lhs) diff --git a/ffcx/codegeneration/numba/__init__.py b/ffcx/codegeneration/numba/__init__.py new file mode 100644 index 000000000..65b8c4f2a --- /dev/null +++ b/ffcx/codegeneration/numba/__init__.py @@ -0,0 +1,5 @@ +"""Generation of numba code.""" + +from ffcx.codegeneration.numba import expressions, file, form, integrals + +__all__ = ["expressions", "file", "form", "integrals", "suffixes"] diff --git a/ffcx/codegeneration/numba/expressions.py b/ffcx/codegeneration/numba/expressions.py new file mode 100644 index 000000000..7908e08a0 --- /dev/null +++ b/ffcx/codegeneration/numba/expressions.py @@ -0,0 +1,110 @@ +# Copyright (C) 2019-2025 Michal Habera, Chris Richardson and Paul T. Kühner +# +# This file is part of FFCx.(https://www.fenicsproject.org) +# +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Generate UFC code for an expression.""" + +import logging +import string + +import numpy as np +import numpy.typing as npt + +from ffcx.codegeneration.backend import FFCXBackend +from ffcx.codegeneration.expression_generator import ExpressionGenerator +from ffcx.codegeneration.numba import expressions_template +from ffcx.codegeneration.numba.implementation import Formatter +from ffcx.ir.representation import ExpressionIR + +logger = logging.getLogger("ffcx") + + +def generator(ir: ExpressionIR, options: dict[str, int | float | npt.DTypeLike]) -> tuple[str, str]: + """Generate UFC code for an expression.""" + logger.info("Generating code for expression:") + assert len(ir.expression.integrand) == 1, "Expressions only support single quadrature rule" + points = next(iter(ir.expression.integrand))[1].points + logger.info(f"--- points: {points}") + factory_name = ir.expression.name + logger.info(f"--- name: {factory_name}") + + # Format declaration + declaration = expressions_template.declaration.format( + factory_name=factory_name, name_from_uflfile=ir.name_from_uflfile + ) + + backend = FFCXBackend(ir, options) + eg = ExpressionGenerator(ir, backend) + + d: dict[str, str | int] = {} + d["name_from_uflfile"] = ir.name_from_uflfile + d["factory_name"] = factory_name + parts = eg.generate() + + # tabulate_expression + num_points = points.shape[0] + num_components = np.prod(ir.expression.shape, dtype=np.int32) + num_argument_dofs = np.prod(ir.expression.tensor_shape, dtype=np.int32) + size_A = num_points * num_components * num_argument_dofs # ref. ufcx.h + size_w = sum(coeff.ufl_element().dim for coeff in ir.expression.coefficient_offsets.keys()) + size_c = sum( + np.prod(constant.ufl_shape, dtype=int) + for constant in ir.expression.original_constant_offsets.keys() + ) + size_coords = ir.expression.number_coordinate_dofs * 3 + size_local_index = 2 # TODO: this is just an upper bound, harmful? + size_permutation = 2 if ir.expression.needs_facet_permutations else 0 + + header = f""" + A = numba.carray(_A, ({size_A})) + w = numba.carray(_w, ({size_w})) + c = numba.carray(_c, ({size_c})) + coordinate_dofs = numba.carray(_coordinate_dofs, ({size_coords})) + entity_local_index = numba.carray(_entity_local_index, ({size_local_index})) + quadrature_permutation = numba.carray(_quadrature_permutation, ({size_permutation})) + """ + F = Formatter(options["scalar_type"]) # type: ignore + body = F.format(parts) + body = "\n".join([" " + line for line in body.split("\n")]) + + d["tabulate_expression"] = header + body + + # TODO: original_coefficient_positions_init + originals = ", ".join(str(i) for i in ir.original_coefficient_positions) + d["original_coefficient_positions"] = f"[{originals}]" + + # TODO: points_init + d["points"] = f"[{', '.join(str(p) for p in points.flatten())}]" + + # TODO: value_shape_init + shape = ", ".join(str(i) for i in ir.expression.shape) + d["value_shape"] = f"[{shape}]" + d["num_components"] = int(num_components) + d["num_coefficients"] = len(ir.expression.coefficient_numbering) + d["num_constants"] = len(ir.constant_names) + d["num_points"] = num_points + d["entity_dimension"] = points.shape[1] + + d["rank"] = len(ir.expression.tensor_shape) + + # TODO: coefficient_names_init + names = ", ".join(f'"{name}"' for name in ir.coefficient_names) + d["coefficient_names"] = f"[{names}]" + + # TODO: constant_names_init + names = ", ".join(f'"{name}"' for name in ir.constant_names) + d["constant_names"] = f"[{names}]" + + d["coordinate_element_hash"] = ir.expression.coordinate_element_hash + + # Check that no keys are redundant or have been missed + 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 + implementation = expressions_template.factory.format_map(d) + + return declaration, implementation diff --git a/ffcx/codegeneration/numba/expressions_template.py b/ffcx/codegeneration/numba/expressions_template.py new file mode 100644 index 000000000..0e8b84b64 --- /dev/null +++ b/ffcx/codegeneration/numba/expressions_template.py @@ -0,0 +1,40 @@ +# Copyright (C) 2019-2025 Michal Habera, Chris Richardson and Paul T.Kühner +# +# This file is part of FFCx.(https://www.fenicsproject.org) +# +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Template for expression output.""" + +declaration = """ +""" + +factory = """ +# Code for expression {factory_name} + +def tabulate_tensor_{factory_name}(_A, _w, _c, _coordinate_dofs, + _entity_local_index, + _quadrature_permutation, custom_data): +{tabulate_expression} + + + +class {factory_name}: + tabulate_tensor = tabulate_tensor_{factory_name} + num_coefficients = {num_coefficients} + num_constants = {num_constants} + original_coefficient_positions = {original_coefficient_positions} + coefficient_names = {coefficient_names} + constant_names = {constant_names} + num_points = {num_points} + entity_dimension = {entity_dimension} + points = {points} + value_shape = {value_shape} + num_components = {num_components} + rank = {rank} + coordinate_element_hash = {coordinate_element_hash} + +# Alias name +{name_from_uflfile} = {factory_name} + +# End of code for expression {factory_name} +""" diff --git a/ffcx/codegeneration/numba/file.py b/ffcx/codegeneration/numba/file.py new file mode 100644 index 000000000..b4a08472e --- /dev/null +++ b/ffcx/codegeneration/numba/file.py @@ -0,0 +1,55 @@ +# Copyright (C) 2009-2025 Anders Logg, Martin Sandve Alnæs, Garth N. Wells, Chris Richardson and +# Paul T. Kühner +# +# This file is part of FFCx.(https://www.fenicsproject.org) +# +# SPDX-License-Identifier: LGPL-3.0-or-later +# +# Note: Most of the code in this file is a direct translation from the +# old implementation in FFC +"""Generate file output for numba.""" + +import logging +import pprint +import textwrap + +from numpy import typing as npt + +from ffcx import __version__ as FFCX_VERSION +from ffcx.codegeneration import __version__ as UFC_VERSION +from ffcx.codegeneration.numba import file_template + +logger = logging.getLogger("ffcx") + + +def generator( + options: dict[str, int | float | npt.DTypeLike], +) -> tuple[tuple[str, str], tuple[str, str]]: + """Generate UFC code for file output. + + Args: + options: Dict of options specified the kenerl generation, these will be documented in the + generated file. + + Returns: tuple of file start- and end sections, each for declaration and implementation. + + """ + logger.info("Generating code for file") + + # Attributes + d = {"ffcx_version": FFCX_VERSION, "ufcx_version": UFC_VERSION} + d["options"] = textwrap.indent(pprint.pformat(options), "# ") + + # Format declaration code + code_pre = ( + file_template.declaration_pre.format_map(d), + file_template.implementation_pre.format_map(d), + ) + + # Format implementation code + code_post = ( + file_template.declaration_post.format_map(d), + file_template.implementation_post.format_map(d), + ) + + return code_pre, code_post diff --git a/ffcx/codegeneration/numba/file_template.py b/ffcx/codegeneration/numba/file_template.py new file mode 100644 index 000000000..0dcbd8ce3 --- /dev/null +++ b/ffcx/codegeneration/numba/file_template.py @@ -0,0 +1,50 @@ +# Copyright (C) 2025 Chris Richardson and Paul T. Kühner +# +# This file is part of FFCx.(https://www.fenicsproject.org) +# +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Template for file output.""" + +declaration_pre = """ +""" + +declaration_post = "" + +implementation_pre = """ +# This code conforms with the UFC specification version {ufcx_version} +# and was automatically generated by FFCx version {ffcx_version}. +# +# WARNING: The numba backend is not production ready. Sub-optimal performance compared to the +# C-backend is to be expected. +# +# This code was generated with the following options: +# +{options} + +import numba +import numpy as np +import math + +# ufcx enums +interval = 10 +triangle = 20 +quadrilateral = 30 +tetrahedron = 40 +hexahedron = 50 +vertex = 60 +prism = 70 +pyramid = 80 + +cell = 0 +exterior_facet = 1 +interior_facet = 2 + +ufcx_basix_element = 0 +ufcx_mixed_element = 1 +ufcx_quadrature_element = 2 +ufcx_basix_custom_element = 3 + +""" + +implementation_post = """ +""" diff --git a/ffcx/codegeneration/numba/form.py b/ffcx/codegeneration/numba/form.py new file mode 100644 index 000000000..f04752352 --- /dev/null +++ b/ffcx/codegeneration/numba/form.py @@ -0,0 +1,154 @@ +# Copyright (C) 2009-2025 Anders Logg, Martin Sandve Alnæs, Chris Richardson and Paul T. Kühner +# +# This file is part of FFCx.(https://www.fenicsproject.org) +# +# SPDX-License-Identifier: LGPL-3.0-or-later + +# Note: Most of the code in this file is a direct translation from the +# old implementation in FFC +"""Template for form output.""" + +import logging +import string + +import numpy as np +from numpy import typing as npt + +from ffcx.codegeneration.numba import form_template + +logger = logging.getLogger("ffcx") + + +def generator(ir, options: dict[str, int | float | npt.DTypeLike]) -> tuple[str, str]: + """Generate UFC code for a form.""" + logger.info("Generating code for form:") + logger.info(f"--- rank: {ir.rank}") + logger.info(f"--- name: {ir.name}") + + d: dict[str, int | str] = {} + d["factory_name"] = ir.name + d["name_from_uflfile"] = ir.name_from_uflfile + d["signature"] = f'"{ir.signature}"' + d["rank"] = ir.rank + d["num_coefficients"] = ir.num_coefficients + + if len(ir.original_coefficient_positions) > 0: + values = ", ".join(str(i) for i in ir.original_coefficient_positions) + + d["original_coefficient_position_init"] = ( + f"original_coefficient_position_{ir.name} = [{values}]" + ) + d["original_coefficient_positions"] = f"original_coefficient_position_{ir.name}" + else: + d["original_coefficient_position_init"] = "" + d["original_coefficient_positions"] = "None" + + if len(ir.coefficient_names) > 0: + values = ", ".join(f'"{name}"' for name in ir.coefficient_names) + d["coefficient_names_init"] = f"coefficient_names_{ir.name} = [{values}]" + d["coefficient_names"] = f"coefficient_names_{ir.name}" + else: + d["coefficient_names_init"] = "" + d["coefficient_names"] = "None" + + d["num_constants"] = ir.num_constants + if ir.num_constants > 0: + d["constant_ranks_init"] = f"constant_ranks_{ir.name} = [{str(ir.constant_ranks)[1:-1]}]" + d["constant_ranks"] = f"constant_ranks_{ir.name}" + + shapes = [ + f"constant_shapes_{ir.name}_{i} = [{str(shape)[1:-1]}]" + for i, shape in enumerate(ir.constant_shapes) + if len(shape) > 0 + ] + names = [f"constant_shapes_{ir.name}_{i}" for i in range(ir.num_constants)] + shapes1 = f"constant_shapes_{ir.name} = [" + for rank, name in zip(ir.constant_ranks, names): + if rank > 0: + shapes1 += f"{name},\n" + else: + shapes1 += "None,\n" + shapes1 += "]" + shapes.append(shapes1) + + d["constant_shapes_init"] = "\n".join(shapes) + d["constant_shapes"] = f"constant_shapes_{ir.name}" + else: + d["constant_ranks_init"] = "" + d["constant_ranks"] = "None" + d["constant_shapes_init"] = "" + d["constant_shapes"] = "None" + + if len(ir.constant_names) > 0: + values = ", ".join(f'"{name}"' for name in ir.constant_names) + d["constant_names_init"] = f"constant_names_{ir.name} = [{values}]" + d["constant_names"] = f"constant_names_{ir.name}" + else: + d["constant_names_init"] = "" + d["constant_names"] = "None" + + if len(ir.finite_element_hashes) > 0: + d["finite_element_hashes"] = f"finite_element_hashes_{ir.name}" + values = ", ".join(f"{0 if el is None else el}" for el in ir.finite_element_hashes) + d["finite_element_hashes_init"] = f"finite_element_hashes_{ir.name} = [{values}]" + else: + d["finite_element_hashes"] = "None" + d["finite_element_hashes_init"] = "" + + integrals = [] + integral_ids = [] + integral_offsets = [0] + integral_domains = [] + # Note: the order of this list is defined by the enum ufcx_integral_type in ufcx.h + for itg_type in ("cell", "exterior_facet", "interior_facet", "vertex", "ridge"): + unsorted_integrals = [] + unsorted_ids = [] + unsorted_domains = [] + for name, domains, id in zip( + ir.integral_names[itg_type], + ir.integral_domains[itg_type], + ir.subdomain_ids[itg_type], + ): + unsorted_integrals += [name] + unsorted_ids += [id] + unsorted_domains += [domains] + + id_sort = np.argsort(unsorted_ids) + integrals += [unsorted_integrals[i] for i in id_sort] + integral_ids += [unsorted_ids[i] for i in id_sort] + integral_domains += [unsorted_domains[i] for i in id_sort] + + integral_offsets.append(sum(len(d) for d in integral_domains)) + + if len(integrals) > 0: + values = ", ".join( + [ + f"{i}_{domain.name}" + for i, domains in zip(integrals, integral_domains) + for domain in domains + ] + ) + d["form_integrals_init"] = f"form_integrals_{ir.name} = [{values}]" + d["form_integrals"] = f"form_integrals_{ir.name}" + values = ", ".join( + f"{i}" for i, domains in zip(integral_ids, integral_domains) for _ in domains + ) + d["form_integral_ids_init"] = f"form_integral_ids_{ir.name} = [{values}]" + d["form_integral_ids"] = f"form_integral_ids_{ir.name}" + else: + d["form_integrals_init"] = "" + d["form_integrals"] = "None" + d["form_integral_ids_init"] = "" + d["form_integral_ids"] = "None" + + values = ", ".join(str(i) for i in integral_offsets) + d["form_integral_offsets_init"] = f"form_integral_offsets_{ir.name} = [{values}]" + + # Check that no keys are redundant or have been missed + fields = [fname for _, fname, _, _ in string.Formatter().parse(form_template.factory) if fname] + assert set(fields) == set(d.keys()), "Mismatch between keys in template and in formatting dict" + + # Format implementation code + implementation = form_template.factory.format_map(d) + + return "", implementation diff --git a/ffcx/codegeneration/numba/form_template.py b/ffcx/codegeneration/numba/form_template.py new file mode 100644 index 000000000..1e4337f3f --- /dev/null +++ b/ffcx/codegeneration/numba/form_template.py @@ -0,0 +1,47 @@ +# Copyright (C) 2025 Chris Richardson +# +# This file is part of FFCx. (https://www.fenicsproject.org) +# +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Template for file output.""" + +factory = """ +# Code for form {factory_name} + +{original_coefficient_position_init} +{finite_element_hashes_init} +{form_integral_offsets_init} +{form_integrals_init} +{form_integral_ids_init} + +{coefficient_names_init} +{constant_names_init} +{constant_ranks_init} +{constant_shapes_init} + +class {factory_name}(object): + + signature = {signature} + rank = {rank} + + num_coefficients = {num_coefficients} + original_coefficient_positions = {original_coefficient_positions} + coefficient_name_map = {coefficient_names} + + num_constants = {num_constants} + constant_ranks = {constant_ranks} + constant_shapes = {constant_shapes} + constant_name_map = {constant_names} + + finite_element_hashes = {finite_element_hashes} + + form_integrals = {form_integrals} + form_integral_ids = {form_integral_ids} + form_integral_offsets = form_integral_offsets_{factory_name} + + +# Alias name +{name_from_uflfile} = {factory_name} + +# End of code for form {factory_name} +""" diff --git a/ffcx/codegeneration/numba/implementation.py b/ffcx/codegeneration/numba/implementation.py new file mode 100644 index 000000000..1afd38486 --- /dev/null +++ b/ffcx/codegeneration/numba/implementation.py @@ -0,0 +1,282 @@ +# Copyright (C) 2025 Chris Richardson and Paul T. Kühner +# +# This file is part of FFCx. (https://www.fenicsproject.org) +# +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Numba implementation for output.""" + +import numpy as np +from numpy import typing as npt + +import ffcx.codegeneration.lnodes as L +from ffcx.codegeneration.utils import dtype_to_scalar_dtype + + +def build_initializer_lists(values: npt.NDArray) -> str: + """Build list of values.""" + arr = "[" + if len(values.shape) == 1: + return "[" + ", ".join(str(v) for v in values) + "]" + elif len(values.shape) > 1: + arr += ",\n".join(build_initializer_lists(v) for v in values) + arr += "]" + return arr + + +class Formatter: + """Implementation for numba output backend.""" + + scalar_type: np.dtype + real_type: np.dtype + + def __init__(self, dtype: npt.DTypeLike) -> None: + """Initialise.""" + self.scalar_type = np.dtype(dtype) + self.real_type = dtype_to_scalar_dtype(dtype) + + def _dtype_to_name(self, dtype: L.DataType) -> str: + """Convert dtype to Python name.""" + if dtype == L.DataType.SCALAR: + return f"np.{self.scalar_type}" + if dtype == L.DataType.REAL: + return f"np.{self.real_type}" + if dtype == L.DataType.INT: + return f"np.{np.int32}" + if dtype == L.DataType.BOOL: + return f"np.{np.bool}" + raise ValueError(f"Invalid dtype: {dtype}") + + def format_section(self, section: L.Section) -> str: + """Format a section.""" + # add new line before section + comments = self._format_comment_str("------------------------") + comments += self._format_comment_str(f"Section: {section.name}") + comments += self._format_comment_str(f"Inputs: {', '.join(w.name for w in section.input)}") + comments += self._format_comment_str( + f"Outputs: {', '.join(w.name for w in section.output)}" + ) + declarations = "".join(self.format(s) for s in section.declarations) + + body = "" + if len(section.statements) > 0: + body = "".join(self.format(s) for s in section.statements) + + body += self._format_comment_str("------------------------") + return comments + declarations + body + + def format_statement_list(self, slist: L.StatementList) -> str: + """Format a list of statements.""" + output = "" + for s in slist.statements: + output += self.format(s) + return output + + def _format_comment_str(self, comment: str) -> str: + """Format str to comment string.""" + return f"# {comment} \n" + + def format_comment(self, c: L.Comment) -> str: + """Format a comment.""" + return self._format_comment_str(c.comment) + + def format_array_decl(self, arr: L.ArrayDecl) -> str: + """Format an array declaration.""" + dtype = arr.symbol.dtype + typename = self._dtype_to_name(dtype) + + symbol = self.format(arr.symbol) + if arr.values is None: + return f"{symbol} = np.empty({arr.sizes}, dtype={typename})\n" + elif arr.values.size == 1: + return f"{symbol} = np.full({arr.sizes}, {arr.values[0]}, dtype={typename})\n" + av = build_initializer_lists(arr.values) + av = f"np.array({av}, dtype={typename})" + return f"{symbol} = {av}\n" + + def format_array_access(self, arr: L.ArrayAccess) -> str: + """Format array access.""" + array = self.format(arr.array) + idx = ", ".join(self.format(ix) for ix in arr.indices) + return f"{array}[{idx}]" + + def format_multi_index(self, index: L.MultiIndex) -> str: + """Format a multi-index.""" + return self.format(index.global_index) + + def format_variable_decl(self, v: L.VariableDecl) -> str: + """Format a variable declaration.""" + sym = self.format(v.symbol) + val = self.format(v.value) + return f"{sym} = {val}\n" + + def format_nary_op(self, oper: L.NaryOp) -> str: + """Format a n argument operation.""" + # Format children + args = [self.format(arg) for arg in oper.args] + + # Apply parentheses + for i in range(len(args)): + if oper.args[i].precedence >= oper.precedence: + args[i] = f"({args[i]})" + + # Return combined string + return f" {oper.op} ".join(args) + + def format_binary_op(self, oper: L.BinOp) -> str: + """Format a binary operation.""" + # Format children + lhs = self.format(oper.lhs) + rhs = self.format(oper.rhs) + + # Apply parentheses + if oper.lhs.precedence >= oper.precedence: + lhs = f"({lhs})" + if oper.rhs.precedence >= oper.precedence: + rhs = f"({rhs})" + + # Return combined string + return f"{lhs} {oper.op} {rhs}" + + def format_neg(self, val: L.Neg) -> str: + """Format unary negation.""" + arg = self.format(val.arg) + return f"-{arg}" + + def format_not(self, val: L.Not) -> str: + """Format not operation.""" + arg = self.format(val.arg) + return f"not({arg})" + + def format_andor(self, oper: L.And | L.Or) -> str: + """Format and or or operation.""" + # Format children + lhs = self.format(oper.lhs) + rhs = self.format(oper.rhs) + + # Apply parentheses + if oper.lhs.precedence >= oper.precedence: + lhs = f"({lhs})" + if oper.rhs.precedence >= oper.precedence: + rhs = f"({rhs})" + + opstr = {"||": "or", "&&": "and"}[oper.op] + + # Return combined string + return f"{lhs} {opstr} {rhs}" + + def format_literal_float(self, val: L.LiteralFloat) -> str: + """Format a literal float.""" + return f"{val.value}" + + def format_literal_int(self, val: L.LiteralInt) -> str: + """Format a literal int.""" + return f"{val.value}" + + def format_for_range(self, r: L.ForRange) -> str: + """Format a loop over a range.""" + begin = self.format(r.begin) + end = self.format(r.end) + index = self.format(r.index) + output = f"for {index} in range({begin}, {end}):\n" + b = self.format(r.body).split("\n") + for line in b: + output += f" {line}\n" + return output + + def format_statement(self, s: L.Statement) -> str: + """Format a statement.""" + return self.format(s.expr) + + def format_assign(self, expr: L.Assign) -> str: + """Format assignment.""" + rhs = self.format(expr.rhs) + lhs = self.format(expr.lhs) + return f"{lhs} {expr.op} {rhs}\n" + + def format_conditional(self, s: L.Conditional) -> str: + """Format a conditional.""" + # Format children + c = self.format(s.condition) + t = self.format(s.true) + f = self.format(s.false) + + # Apply parentheses + if s.condition.precedence >= s.precedence: + c = f"({c})" + if s.true.precedence >= s.precedence: + t = f"({t})" + if s.false.precedence >= s.precedence: + f = f"({f})" + + # Return combined string + return f"({t} if {c} else {f})" + + def format_symbol(self, s: L.Symbol) -> str: + """Format a symbol.""" + return f"{s.name}" + + def format_mathfunction(self, f: L.MathFunction) -> str: + """Format a math function.""" + function_map = { + "ln": "log", + "acos": "arccos", + "asin": "arcsin", + "atan": "arctan", + "atan2": "arctan2", + "acosh": "arccosh", + "asinh": "arcsinh", + "atanh": "arctanh", + } + function = function_map.get(f.function, f.function) + args = [self.format(arg) for arg in f.args] + if "bessel_y" in function: + return "scipy.special.yn" + if "bessel_j" in function: + return "scipy.special.jn" + if function == "erf": + return f"math.erf({args[0]})" + argstr = ", ".join(args) + return f"np.{function}({argstr})" + + impl = { + "StatementList": format_statement_list, + "Comment": format_comment, + "Section": format_section, + "ArrayDecl": format_array_decl, + "ArrayAccess": format_array_access, + "MultiIndex": format_multi_index, + "VariableDecl": format_variable_decl, + "ForRange": format_for_range, + "Statement": format_statement, + "Assign": format_assign, + "AssignAdd": format_assign, + "Product": format_nary_op, + "Sum": format_nary_op, + "Add": format_binary_op, + "Sub": format_binary_op, + "Mul": format_binary_op, + "Div": format_binary_op, + "Neg": format_neg, + "Not": format_not, + "LiteralFloat": format_literal_float, + "LiteralInt": format_literal_int, + "Symbol": format_symbol, + "Conditional": format_conditional, + "MathFunction": format_mathfunction, + "And": format_andor, + "Or": format_andor, + "NE": format_binary_op, + "EQ": format_binary_op, + "GE": format_binary_op, + "LE": format_binary_op, + "GT": format_binary_op, + "LT": format_binary_op, + } + + def format(self, s: L.LNode) -> str: + """Format output.""" + name = s.__class__.__name__ + try: + return self.impl[name](self, s) # type: ignore + except KeyError: + raise RuntimeError("Unknown statement: ", name) diff --git a/ffcx/codegeneration/numba/integrals.py b/ffcx/codegeneration/numba/integrals.py new file mode 100644 index 000000000..c72224c68 --- /dev/null +++ b/ffcx/codegeneration/numba/integrals.py @@ -0,0 +1,107 @@ +# Copyright (C) 2015-2021 Martin Sandve Alnæs, Michal Habera, Igor Baratta, Chris Richardson and +# Paul T. Kühner +# +# This file is part of FFCx. (https://www.fenicsproject.org) +# +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Generate UFC code for an integral.""" + +import logging +import string + +import basix +import numpy as np +from numpy import typing as npt + +from ffcx.codegeneration.backend import FFCXBackend +from ffcx.codegeneration.integral_generator import IntegralGenerator +from ffcx.codegeneration.numba import integrals_template as ufcx_integrals +from ffcx.codegeneration.numba.implementation import Formatter +from ffcx.ir.representation import IntegralIR + +logger = logging.getLogger("ffcx") + + +def generator( + ir: IntegralIR, domain: basix.CellType, options: dict[str, int | float | npt.DTypeLike] +) -> tuple[str, str]: + """Generate numba code for an integral. + + Args: + ir: IR of the integral + domain: basix cell type + options: dict of kernel generation options + + Returns: + Tuple of declaration (header) and implementation (source) strings. + + Note: + numba backend only provides a declaration. Implementation string will always be empty. + + """ + logger.info("Generating code for integral:") + logger.info(f"--- type: {ir.expression.integral_type}") + logger.info(f"--- name: {ir.expression.name}") + + factory_name = f"{ir.expression.name}_{domain.name}" + + # Format declaration + declaration = "" + + # Create FFCx backend + backend = FFCXBackend(ir, options) + + # Configure kernel generator + ig = IntegralGenerator(ir, backend) + + # Generate code AST for the tabulate_tensor body + parts = ig.generate(domain) + + # Format code as string + F = Formatter(options["scalar_type"]) # type: ignore + body = F.format(parts) + body = "\n".join([" " + line for line in body.split("\n")]) + + # Generate generic FFCx code snippets and add specific parts + d: dict[str, str] = {} + + d["factory_name"] = factory_name + + # TODO: enabled_coefficients_init - required? + vals = ", ".join("1" if i else "0" for i in ir.enabled_coefficients) + d["enabled_coefficients"] = f"[{vals}]" + + # tabulate_tensor + # Note: In contrast to the C implementation we actually need to provide/compute the sizes of the + # array. + size_A = np.prod(ir.expression.tensor_shape) + size_w = sum(coeff.ufl_element().dim for coeff in ir.expression.coefficient_offsets.keys()) + size_c = sum( + np.prod(constant.ufl_shape, dtype=int) + for constant in ir.expression.original_constant_offsets.keys() + ) + size_coords = ir.expression.number_coordinate_dofs * 3 + size_local_index = 2 # TODO: this is just an upper bound + size_permutation = 2 if ir.expression.needs_facet_permutations else 0 + + header = f""" + A = numba.carray(_A, ({size_A})) + w = numba.carray(_w, ({size_w})) + c = numba.carray(_c, ({size_c})) + coordinate_dofs = numba.carray(_coordinate_dofs, ({size_coords})) + entity_local_index = numba.carray(_entity_local_index, ({size_local_index})) + quadrature_permutation = numba.carray(_quadrature_permutation, ({size_permutation})) + """ + d["tabulate_tensor"] = header + body + d["needs_facet_permutations"] = "True" if ir.expression.needs_facet_permutations else "False" + d["coordinate_element_hash"] = ir.expression.coordinate_element_hash + d["domain"] = str(int(domain)) + + assert ir.expression.coordinate_element_hash is not None + implementation = ufcx_integrals.factory.format_map(d) + + # Check that no keys are redundant or have been missed + fields = [fname for _, fname, _, _ in string.Formatter().parse(ufcx_integrals.factory) if fname] + assert set(fields) == set(d.keys()), "Mismatch between keys in template and in formatting dict" + + return declaration, implementation diff --git a/ffcx/codegeneration/numba/integrals_template.py b/ffcx/codegeneration/numba/integrals_template.py new file mode 100644 index 000000000..28dccd0ea --- /dev/null +++ b/ffcx/codegeneration/numba/integrals_template.py @@ -0,0 +1,23 @@ +# Copyright (C) 2025 Chris Richardson and Paul T. Kühner +# +# This file is part of FFCx. (https://www.fenicsproject.org) +# +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Template for integral output.""" + +factory = """ +# Code for integral {factory_name} + +def tabulate_tensor_{factory_name}(_A, _w, _c, _coordinate_dofs, + _entity_local_index, _quadrature_permutation, custom_data): +{tabulate_tensor} + +class {factory_name}(object): + enabled_coefficients = {enabled_coefficients} + tabulate_tensor = tabulate_tensor_{factory_name} + needs_facet_permutations = {needs_facet_permutations} + coordinate_element_hash = {coordinate_element_hash} + domain = {domain} + +# End of code for integral {factory_name} +""" diff --git a/ffcx/compiler.py b/ffcx/compiler.py index aaccd4d24..61f94e61b 100644 --- a/ffcx/compiler.py +++ b/ffcx/compiler.py @@ -93,12 +93,15 @@ def compile_ufl_objects( """Generate UFC code for a given UFL objects. Args: - ufl_objects: Objects to be compiled. Accepts elements, forms, - integrals or coordinate mappings. - object_names: Map from object Python id to object name - prefix: Prefix - options: Options - visualise: Toggle visualisation + ufl_objects: Objects to be compiled. Accepts elements, forms, + integrals or coordinate mappings. + object_names: Map from object Python id to object name + prefix: Prefix + options: Options + visualise: Toggle visualisation + + Returns: tuple of declaration and implementation. + """ _object_names = object_names if object_names is not None else {} _prefix = prefix if prefix is not None else "" diff --git a/ffcx/formatting.py b/ffcx/formatting.py index 553179ca5..5396c7417 100644 --- a/ffcx/formatting.py +++ b/ffcx/formatting.py @@ -16,7 +16,7 @@ from __future__ import annotations import logging -import os +from pathlib import Path from ffcx.codegeneration.codegeneration import CodeBlocks @@ -38,14 +38,17 @@ def format_code(code: CodeBlocks) -> tuple[str, str]: return code_h, code_c -def write_code(code_h: str, code_c: str, prefix: str, output_dir: str) -> None: +def write_code( + code_h: str, code_c: str, prefix: str, suffixes: tuple[str | None, str | None], output_dir: str +) -> None: """Write code to files.""" - _write_file(code_h, prefix, ".h", output_dir) - _write_file(code_c, prefix, ".c", output_dir) + if suffixes[0] is not None: + _write_file(code_h, prefix, suffixes[0], output_dir) + if suffixes[1] is not None: + _write_file(code_c, prefix, suffixes[1], output_dir) -def _write_file(output: str, prefix: str, postfix: str, output_dir: str) -> None: +def _write_file(output: str, prefix: str, suffix: str, output_dir: str) -> None: """Write generated code to file.""" - filename = os.path.join(output_dir, prefix + postfix) - with open(filename, "w") as hfile: - hfile.write(output) + with open(Path(output_dir) / (prefix + suffix), "w") as file: + file.write(output) diff --git a/ffcx/ir/integral.py b/ffcx/ir/integral.py index aaed1e63b..a89a91082 100644 --- a/ffcx/ir/integral.py +++ b/ffcx/ir/integral.py @@ -80,6 +80,7 @@ class CommonExpressionIR(typing.NamedTuple): needs_facet_permutations: bool shape: list[int] coordinate_element_hash: str + number_coordinate_dofs: int class ModifiedArgumentDataT(typing.NamedTuple): diff --git a/ffcx/ir/representation.py b/ffcx/ir/representation.py index 627db639d..b0c7fdf00 100644 --- a/ffcx/ir/representation.py +++ b/ffcx/ir/representation.py @@ -234,6 +234,7 @@ def _compute_integral_ir( "entity_type": entity_type, "shape": (), "coordinate_element_hash": itg_data.domain.ufl_coordinate_element().basix_hash(), + "number_coordinate_dofs": itg_data.domain.ufl_coordinate_element().dim, } ir = { "rank": form_data.rank, @@ -613,6 +614,9 @@ def _compute_expression_ir( base_ir["coordinate_element_hash"] = ( expr_domain.ufl_coordinate_element().basix_hash() if expr_domain is not None else 0 ) + base_ir["number_coordinate_dofs"] = ( + 0 if expr_domain is None else expr_domain.ufl_coordinate_element().dim + ) weights = np.array([1.0] * points.shape[0]) rule = QuadratureRule(points, weights) diff --git a/ffcx/main.py b/ffcx/main.py index 84978fa00..0a72e2eaf 100644 --- a/ffcx/main.py +++ b/ffcx/main.py @@ -83,8 +83,16 @@ def main(args: Sequence[str] | None = None) -> int: visualise=xargs.visualise, ) + # File suffixes + # TODO: this needs to be moved into the language backends + suffixes: tuple[str | None, str | None] + if options["language"] == "C": + suffixes = (".h", ".c") + else: # numba: + suffixes = (None, "_numba.py") + # Write to file - formatting.write_code(code_h, code_c, prefix, xargs.output_directory) + formatting.write_code(code_h, code_c, prefix, suffixes, xargs.output_directory) # Turn off profiling and write status to file if xargs.profile: diff --git a/ffcx/options.py b/ffcx/options.py index cd6ec8d1a..d9908cc88 100644 --- a/ffcx/options.py +++ b/ffcx/options.py @@ -20,6 +20,7 @@ logger = logging.getLogger("ffcx") FFCX_DEFAULT_OPTIONS = { + "language": (str, "C", "language to output kernel in", ("C", "numba")), "epsilon": (float, 1e-14, "machine precision, used for dropping zero terms in tables.", None), "scalar_type": ( str, diff --git a/pyproject.toml b/pyproject.toml index 4793872a1..bff7afc91 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,6 +52,7 @@ packages = [ "ffcx", "ffcx.codegeneration", "ffcx.codegeneration.C", + "ffcx.codegeneration.numba", "ffcx.ir", "ffcx.ir.analysis", ] diff --git a/test/Poisson.py b/test/poisson.py similarity index 69% rename from test/Poisson.py rename to test/poisson.py index b9c1ba503..c7468a5c7 100644 --- a/test/Poisson.py +++ b/test/poisson.py @@ -1,4 +1,4 @@ -# Copyright (C) 2004-2007 Anders Logg +# Copyright (C) 2004-2025 Anders Logg and Paul T. Kühner # # This file is part of FFCx. # @@ -31,19 +31,27 @@ dx, grad, inner, + tr, ) -mesh = Mesh(basix.ufl.element("P", "triangle", 2, shape=(2,))) +mesh = Mesh(basix.ufl.element("P", "triangle", 1, shape=(2,))) -e = basix.ufl.element("Lagrange", "triangle", 2) +# Forms +e = basix.ufl.element("Lagrange", "triangle", 1) space = FunctionSpace(mesh, e) u = TrialFunction(space) v = TestFunction(space) f = Coefficient(space) -kappa1 = Constant(mesh, shape=(2, 2)) -kappa2 = Constant(mesh, shape=(2, 2)) +kappa = Constant(mesh, shape=(2, 2)) -a = inner(kappa1, kappa2) * inner(grad(u), grad(v)) * dx +a = tr(kappa) * inner(grad(u), grad(v)) * dx L = f * v * dx + +# Expressions +e_vec = basix.ufl.element("Lagrange", "triangle", 1, shape=(2, 3)) +space_vec = FunctionSpace(mesh, e_vec) +f_vec = Coefficient(space_vec) + +expressions = [(kappa * f_vec, e_vec.basix_element.points)] diff --git a/test/test_cmdline.py b/test/test_cmdline.py index 35b29b47c..1c36930e7 100644 --- a/test/test_cmdline.py +++ b/test/test_cmdline.py @@ -1,28 +1,26 @@ -# Copyright (C) 2018 Chris N. Richardson +# Copyright (C) 2018-2025 Chris N. Richardson and Paul T. Kühner # # This file is part of FFCx. (https://www.fenicsproject.org) # # SPDX-License-Identifier: LGPL-3.0-or-later -import os -import os.path +import importlib import subprocess +from pathlib import Path import pytest def test_cmdline_simple(): - os.chdir(os.path.dirname(__file__)) - subprocess.run(["ffcx", "Poisson.py"], check=True) + dir = Path(__file__).parent + subprocess.run(["ffcx", "-o", dir, dir / "poisson.py"], check=True) +@pytest.mark.skipif( + importlib.util.find_spec("pygraphviz") is None, reason="pygraphviz not installed." +) def test_visualise(): - try: - import pygraphviz # noqa: F401 - except ImportError: - pytest.skip("pygraphviz not installed") - - os.chdir(os.path.dirname(__file__)) - subprocess.run(["ffcx", "--visualise", "Poisson.py"]) - assert os.path.isfile("S.pdf") - assert os.path.isfile("F.pdf") + dir = Path(__file__).parent + subprocess.run(["ffcx", "-o", dir, "--visualise", dir / "poisson.py"], check=True) + assert (Path.cwd() / "S.pdf").is_file() + assert (Path.cwd() / "F.pdf").is_file() diff --git a/test/test_lnodes.py b/test/test_lnodes.py index 1be46d8d0..e8187bbd9 100644 --- a/test/test_lnodes.py +++ b/test/test_lnodes.py @@ -5,7 +5,7 @@ from cffi import FFI from ffcx.codegeneration import lnodes as L -from ffcx.codegeneration.C.c_implementation import CFormatter +from ffcx.codegeneration.C.implementation import Formatter from ffcx.codegeneration.utils import dtype_to_c_type @@ -32,10 +32,10 @@ def test_gemm(dtype): code += [L.ForRange(k, 0, r, body=body)] # Format into C and compile with CFFI - Q = CFormatter(dtype=dtype) + Q = Formatter(dtype=dtype) c_scalar = dtype_to_c_type(dtype) decl = f"void gemm({c_scalar} *A, {c_scalar} *B, {c_scalar} *C)" - c_code = decl + "{\n" + Q.c_format(L.StatementList(code)) + "\n}\n" + c_code = decl + "{\n" + Q.format(L.StatementList(code)) + "\n}\n" ffibuilder = FFI() ffibuilder.cdef(decl + ";") @@ -75,10 +75,10 @@ def test_gemv(dtype): code += [L.ForRange(j, 0, q, body=body)] # Format into C and compile with CFFI - Q = CFormatter(dtype=dtype) + Q = Formatter(dtype=dtype) c_scalar = dtype_to_c_type(dtype) decl = f"void gemm({c_scalar} *y, {c_scalar} *A, {c_scalar} *x)" - c_code = decl + "{\n" + Q.c_format(L.StatementList(code)) + "\n}\n" + c_code = decl + "{\n" + Q.format(L.StatementList(code)) + "\n}\n" ffibuilder = FFI() ffibuilder.cdef(decl + ";") diff --git a/test/test_numba.py b/test/test_numba.py new file mode 100644 index 000000000..335b5f0e6 --- /dev/null +++ b/test/test_numba.py @@ -0,0 +1,121 @@ +# Copyright (C) 2025 Paul T. Kühner +# +# This file is part of FFCx.(https://www.fenicsproject.org) +# +# SPDX-License-Identifier: LGPL-3.0-or-later + +import ctypes +import importlib +from pathlib import Path + +import numba +import numpy as np +import numpy.typing as npt +import pytest + +import ffcx.main +from ffcx.codegeneration.utils import dtype_to_scalar_dtype, numba_ufcx_kernel_signature + + +def wrap_kernel(scalar_type, real_type): + c_signature = numba_ufcx_kernel_signature(scalar_type, real_type) + return numba.cfunc(c_signature, nopython=True) + + +def as_C_array(np_array: npt.NDArray): + dtype_C = np.ctypeslib.as_ctypes_type(np_array.dtype) + return np_array.ctypes.data_as(ctypes.POINTER(dtype_C)) + + +@pytest.mark.parametrize("scalar_type", ["float32", "float64"]) # TODO: complex limited by ctypes +def test_integral(scalar_type: str) -> None: + opts = f"--language numba --scalar_type {scalar_type}" + dir = Path(__file__).parent + assert ffcx.main.main([str(dir / "poisson.py"), *opts.split(" ")]) == 0 + + poisson = importlib.import_module("poisson_numba") + + dtype = np.dtype(scalar_type).type + dtype_r = dtype_to_scalar_dtype(dtype) + + kernel_a = wrap_kernel(dtype, dtype_r)(poisson.form_poisson_a.form_integrals[0].tabulate_tensor) + + A = np.zeros((3, 3), dtype=dtype) + w = np.array([], dtype=dtype) + kappa_value = np.array([[1.0, 2.0], [3.0, 4.0]]) + c = np.array(kappa_value.flatten(), dtype=dtype) + coords = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=dtype_r) + empty = np.empty((0,), dtype=dtype_r) + + kernel_a( + as_C_array(A), + as_C_array(w), + as_C_array(c), + as_C_array(coords), + as_C_array(empty), + as_C_array(empty), + 0, + ) + A_expected = np.array( + [[1.0, -0.5, -0.5], [-0.5, 0.5, 0.0], [-0.5, 0.0, 0.5]], dtype=scalar_type + ) + + assert np.allclose(A, np.trace(kappa_value) * A_expected) + + kernel_L = wrap_kernel(dtype, dtype_r)(poisson.form_poisson_L.form_integrals[0].tabulate_tensor) + + b = np.zeros((3,), dtype=dtype) + w = np.full((3,), 0.5, dtype=dtype) + c = np.empty((0,), dtype=dtype) + coords = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=dtype_r) + empty = np.empty((0,), dtype=dtype_r) + + kernel_L( + as_C_array(b), + as_C_array(w), + as_C_array(c), + as_C_array(coords), + as_C_array(empty), + as_C_array(empty), + 0, + ) + + b_expected = np.full((3,), 1 / 6, dtype=np.float64) + assert np.allclose(b, 0.5 * b_expected) + + +@pytest.mark.parametrize("scalar_type", ["float32", "float64"]) # TODO: complex limited by ctypes +def test_expression(scalar_type: str) -> None: + opts = f"--language numba --scalar_type {scalar_type}" + dir = Path(__file__).parent + assert ffcx.main.main([str(dir / "poisson.py"), *opts.split(" ")]) == 0 + + poisson = importlib.import_module("poisson_numba") + + dtype = np.dtype(scalar_type).type + dtype_r = dtype_to_scalar_dtype(dtype) + + kernel_expr = wrap_kernel(dtype, dtype_r)(poisson.expression_poisson_0.tabulate_tensor) + + e = np.zeros((6 * 3,), dtype=dtype) + w = np.array( + [1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9, 10, 10, 11, 11], dtype=dtype + ) + kappa_value = np.array([[1.0, 2.0], [3.0, 4.0]]) + c = np.array(kappa_value.flatten(), dtype=dtype) + coords = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=dtype_r) + empty = np.empty((0,), dtype=dtype_r) + + kernel_expr( + as_C_array(e), + as_C_array(w), + as_C_array(c), + as_C_array(coords), + as_C_array(empty), + as_C_array(empty), + 0, + ) + e_expected = np.array( + [5, 7, 8, 11, 15, 18, 14, 16, 17, 32, 36, 39, 23, 25, 26, 53, 57, 60], dtype=dtype + ) + assert np.allclose(e, e_expected)