diff --git a/ffcx_backends/cpp.py b/ffcx_backends/cpp.py index 6283ddc..1631ae4 100644 --- a/ffcx_backends/cpp.py +++ b/ffcx_backends/cpp.py @@ -4,7 +4,7 @@ import logging import pprint import textwrap -from typing import Any, ClassVar +from typing import ClassVar import basix import ffcx.codegeneration.lnodes as L # noqa @@ -19,20 +19,6 @@ logger = logging.getLogger("ffcx") -def dtype_to_cpp_type(dtype: L.DataType, scalar_type: str, real_type: str) -> str: - """Map L.DataType to C++ type.""" - if dtype == L.DataType.SCALAR: - return scalar_type - elif dtype == L.DataType.REAL: - return real_type - elif dtype == L.DataType.INT: - return "std::int32_t" - elif dtype == L.DataType.BOOL: - return "bool" - else: - raise ValueError(f"Invalid datatype: {dtype}") - - class Formatter: """Format FFCx nodes into C++.""" @@ -76,10 +62,29 @@ def build_initializer_lists(values: npt.NDArray) -> str: arr += "}" return arr - def __init__(self, scalar: Any) -> None: - """Initialise.""" - self.scalar_type = "T" - self.real_type = "U" + def __init__(self, scalar_geometry: bool = False) -> None: + """Initialise. + + Args: + scalar_geometry: Force scalar type == geometry type in the kernel body. + """ + self._scalar_type = "T" + self._geometry_type = "U" + if scalar_geometry: + self._geometry_type = self._scalar_type + + def dtype_to_cpp_type(self, dtype: L.DataType) -> str: + """Map L.DataType to C++ type.""" + if dtype == L.DataType.SCALAR: + return self._scalar_type + elif dtype == L.DataType.REAL: + return self._geometry_type + elif dtype == L.DataType.INT: + return "std::int32_t" + elif dtype == L.DataType.BOOL: + return "bool" + else: + raise ValueError(f"Invalid datatype: {dtype}") @functools.singledispatchmethod def __call__(self, obj: L.LNode) -> str: @@ -122,7 +127,7 @@ def format_array_decl(self, arr: L.ArrayDecl) -> str: dtype = arr.symbol.dtype assert dtype is not None - typename = dtype_to_cpp_type(dtype, self.scalar_type, self.real_type) + typename = self.dtype_to_cpp_type(dtype) symbol = self(arr.symbol) dims = "".join([f"[{i}]" for i in arr.sizes]) @@ -152,7 +157,7 @@ def format_variable_decl(self, v: L.VariableDecl) -> str: val = self(v.value) symbol = self(v.symbol) assert v.symbol.dtype - typename = dtype_to_cpp_type(v.symbol.dtype, self.scalar_type, self.real_type) + typename = self.dtype_to_cpp_type(v.symbol.dtype) return f"{typename} {symbol} = {val};\n" @__call__.register @@ -329,7 +334,8 @@ def generator(ir: ExpressionIR, options: dict[str, int | float | npt.DTypeLike]) d["factory_name"] = factory_name parts = eg.generate() - formatter = Formatter(options["scalar_type"]) + scalar_geometry = True if options.get("scalar_geometry") else False + formatter = Formatter(scalar_geometry) d["tabulate_expression"] = formatter(parts) if len(ir.original_coefficient_positions) > 0: @@ -442,7 +448,8 @@ def generator( parts = ig.generate(domain) # Format code as string - formatter = Formatter(options["scalar_type"]) + scalar_geometry = True if options.get("scalar_geometry") else False + formatter = Formatter(scalar_geometry) body = formatter(parts) # Generate generic FFCx code snippets and add specific parts diff --git a/test/test_cpp.py b/test/test_cpp.py index 57cd16c..1215776 100644 --- a/test/test_cpp.py +++ b/test/test_cpp.py @@ -1,8 +1,14 @@ +import re + import basix.ufl +import ffcx.codegeneration.lnodes as L # noqa: N812 +import pytest import ufl from ffcx.compiler import compile_ufl_objects from ffcx.options import get_options +from ffcx_backends.cpp import Formatter + def test_integral() -> None: element = basix.ufl.element("Lagrange", "triangle", 1) @@ -16,3 +22,52 @@ def test_integral() -> None: compiled_objects = compile_ufl_objects([a], opts) assert len(compiled_objects) == 2 + + +@pytest.mark.parametrize(("scalar_geometry", "expected"), [(False, "U"), (True, "T"), (None, "U")]) +def test_scalar_geometry(scalar_geometry: bool | None, expected: str) -> None: + """Geometry-derived temporaries are emitted in T only with ``scalar_geometry``. + + The kernel signature is unaffected: coordinate dofs are always read as ``U``, + the option only changes the type they are computed in. + """ + element = basix.ufl.element("Lagrange", "triangle", 1) + domain = ufl.Mesh(basix.ufl.element("Lagrange", "triangle", 1, shape=(2,))) + space = ufl.FunctionSpace(domain, element) + u, v = ufl.TrialFunction(space), ufl.TestFunction(space) + + a = (ufl.inner(u, v) + ufl.inner(ufl.grad(u), ufl.grad(v))) * ufl.dx + opts = get_options({"language": "ffcx_backends.cpp"}) + if scalar_geometry is not None: + opts["scalar_geometry"] = scalar_geometry + + code = compile_ufl_objects([a], opts)[0][0] + + # Quadrature weights, tabulated basis functions and the Jacobian are all + # geometry-derived, hence declared in the geometry type. + for declaration in [ + r"static const (\w+) weights_\w+\[", + r"static const (\w+) FE\w+\[", + # The Jacobian symbol carries a process-global counter, hence J. + r"^(\w+) J\d+_c0 = ", + ]: + matches = re.findall(declaration, code, flags=re.MULTILINE) + assert matches, f"no declaration matching {declaration} in generated code" + assert set(matches) == {expected} + + # The tabulate_tensor signature keeps the scalar/geometry split either way. + signature = re.search(r"static void tabulate_tensor\((.*?)\)\s*\{", code, flags=re.DOTALL) + assert signature is not None + assert "const U* RESTRICT coordinate_dofs" in signature.group(1) + assert "T* RESTRICT A" in signature.group(1) + + +@pytest.mark.parametrize(("scalar_geometry", "expected"), [(False, "U"), (True, "T")]) +def test_formatter_dtype_to_cpp_type(scalar_geometry: bool, expected: str) -> None: + """REAL follows the geometry type, the remaining datatypes do not.""" + formatter = Formatter(scalar_geometry) + + assert formatter.dtype_to_cpp_type(L.DataType.REAL) == expected + assert formatter.dtype_to_cpp_type(L.DataType.SCALAR) == "T" + assert formatter.dtype_to_cpp_type(L.DataType.INT) == "std::int32_t" + assert formatter.dtype_to_cpp_type(L.DataType.BOOL) == "bool"