Add scalar geometry option

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

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<n>.
+        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"