Support ufl.interpolate

ffcx/codegeneration/C/formatter.py

@@ -239,6 +239,12 @@ def _(self, arr: L.ArrayDecl) -> str:
         cstr = "static const " if arr.const else ""
         return f"{cstr}{typename} {symbol}{dims} = {vals};\n"
 
+    @__call__.register
+    def _(self, arr: L.CallOp) -> str:
+        """Format a function call."""
+        args = ", ".join(self(arg) for arg in arr.args)
+        return f"{arr.func}_{self.scalar_type}({args});\n"
+
     @__call__.register
     def _(self, arr: L.ArrayAccess) -> str:
         """Format an array access."""

ffcx/codegeneration/access.py

@@ -12,6 +12,7 @@
 import ufl
 
 import ffcx.codegeneration.lnodes as L
+from ffcx.analysis import ProxyCoefficient
 from ffcx.definitions import entity_types
 from ffcx.ir.analysis.modified_terminals import ModifiedTerminal
 from ffcx.ir.elementtables import UniqueTableReferenceT
@@ -36,6 +37,7 @@ def __init__(self, entity_type: entity_types, integral_type: str, symbols, optio
         # Lookup table for handler to call when the "get" method (below)
         # is called, depending on the first argument type.
         self.call_lookup = {
+            ProxyCoefficient: self.proxy_coefficient,
             ufl.coefficient.Coefficient: self.coefficient,
             ufl.constant.Constant: self.constant,
             ufl.geometry.Jacobian: self.jacobian,
@@ -78,6 +80,30 @@ def get(
         else:
             raise RuntimeError(f"Not handled: {type(e)}")
 
+    def proxy_coefficient(
+        self,
+        mt: ModifiedTerminal,
+        tabledata: UniqueTableReferenceT,
+        quadrature_rule: QuadratureRule,
+    ):
+        """Access a coefficient."""
+        ttype = tabledata.ttype
+        assert ttype != "zeros"
+
+        num_dofs = tabledata.values.shape[3]
+        begin = tabledata.offset
+        assert begin is not None
+        assert tabledata.block_size is not None
+        end = begin + tabledata.block_size * (num_dofs - 1) + 1
+        if ttype == "ones" and (end - begin) == 1:
+            # f = 1.0 * f_{begin}, just return direct reference to dof
+            # array at dof begin (if mt is restricted, begin contains
+            # cell offset)
+            return self.symbols.proxy_coefficient_dof_access(mt.terminal, begin)
+        else:
+            # Return symbol, see definitions for computation
+            return self.symbols.proxy_coefficient_value(mt)
+
     def coefficient(
         self,
         mt: ModifiedTerminal,

ffcx/codegeneration/backend.py

@@ -20,11 +20,16 @@ def __init__(self, ir: IntegralIR | ExpressionIR, options):
         """Initialise."""
         coefficient_numbering = ir.expression.coefficient_numbering
         coefficient_offsets = ir.expression.coefficient_offsets
-
+        proxy_coefficient_numbering = ir.expression.proxy_coefficient_numbering
+        proxy_coefficient_offsets = ir.expression.proxy_coefficient_offsets
         original_constant_offsets = ir.expression.original_constant_offsets
 
         self.symbols = FFCXBackendSymbols(
-            coefficient_numbering, coefficient_offsets, original_constant_offsets
+            coefficient_numbering,
+            coefficient_offsets,
+            original_constant_offsets,
+            proxy_coefficient_numbering,
+            proxy_coefficient_offsets,
         )
         self.access = FFCXBackendAccess(
             ir.expression.entity_type, ir.expression.integral_type, self.symbols, options

ffcx/codegeneration/codegeneration.py

@@ -32,9 +32,9 @@ class CodeBlocks(typing.NamedTuple):
     """
 
     file_pre: list[tuple[str, str]]
+    expressions: list[tuple[str, str]]
     integrals: list[tuple[str, str]]
     forms: list[tuple[str, str]]
-    expressions: list[tuple[str, str]]
     file_post: list[tuple[str, str]]
 
 

ffcx/codegeneration/definitions.py

@@ -6,10 +6,12 @@
 """FFCx/UFCx specific variable definitions."""
 
 import logging
+import typing
 
 import ufl
 
 import ffcx.codegeneration.lnodes as L
+from ffcx.analysis import ProxyCoefficient
 from ffcx.definitions import entity_types
 from ffcx.ir.analysis.modified_terminals import ModifiedTerminal
 from ffcx.ir.elementtables import UniqueTableReferenceT
@@ -51,6 +53,7 @@ class FFCXBackendDefinitions:
     """FFCx specific code definitions."""
 
     entity_type: entity_types
+    handler_lookup: dict[ufl.core.ufl_type.UFLType, typing.Callable]
 
     def __init__(self, entity_type: entity_types, integral_type: str, access, options):
         """Initialise."""
@@ -63,6 +66,7 @@ def __init__(self, entity_type: entity_types, integral_type: str, access, option
         # called, depending on the first argument type.
         self.handler_lookup = {
             ufl.coefficient.Coefficient: self.coefficient,
+            ProxyCoefficient: self.proxy_coefficient,
             ufl.geometry.Jacobian: self._define_coordinate_dofs_lincomb,
             ufl.geometry.SpatialCoordinate: self.spatial_coordinate,
             ufl.constant.Constant: self.pass_through,
@@ -167,7 +171,65 @@ def coefficient(
         # assert input and output are Symbol objects
         assert all(isinstance(i, L.Symbol) for i in input)
         assert all(isinstance(o, L.Symbol) for o in output)
+        return L.Section(name, code, declaration, input, output, annotations)
+
+    def proxy_coefficient(
+        self,
+        mt: ModifiedTerminal,
+        tabledata: UniqueTableReferenceT,
+        quadrature_rule: QuadratureRule,
+        access: L.Symbol,
+    ) -> L.Section | list:
+        """Return definition code for coefficients."""
+        # For applying tensor product to coefficients, we need to know
+        # if the coefficient has a tensor factorisation and if the
+        # quadrature rule has a tensor factorisation. If both are true,
+        # we can apply the tensor product to the coefficient.
+
+        iq_symbol = self.symbols.quadrature_loop_index
+        ic_symbol = self.symbols.coefficient_dof_sum_index
+
+        iq = create_quadrature_index(quadrature_rule, iq_symbol)
+        ic = create_dof_index(tabledata, ic_symbol)
+
+        # Get properties of tables
+        ttype = tabledata.ttype
+        num_dofs = tabledata.values.shape[3]
+        bs = tabledata.block_size
+        begin = tabledata.offset
+        assert bs is not None
+        assert begin is not None
+        end = begin + bs * (num_dofs - 1) + 1
 
+        if ttype == "zeros":
+            logger.debug("Not expecting zero coefficients to get this far.")
+            return []
+
+        # For a constant coefficient we reference the dofs directly, so
+        # no definition needed
+        if ttype == "ones" and end - begin == 1:
+            return []
+
+        assert begin < end
+
+        # Get access to element table
+        FE, tables = self.access.table_access(tabledata, self.entity_type, mt.restriction, iq, ic)
+        dof_access: L.ArrayAccess = self.symbols.proxy_coefficient_dof_access(
+            mt.terminal, (ic.global_index) * bs + begin
+        )
+
+        declaration: list[L.Declaration] = [L.VariableDecl(access, 0.0)]
+        body = [L.AssignAdd(access, dof_access * FE)]
+        code = [L.create_nested_for_loops([ic], body)]
+
+        name = type(mt.terminal).__name__
+        input = [dof_access.array, *tables]
+        output = [access]
+        annotations = [L.Annotation.fuse]
+
+        # assert input and output are Symbol objects
+        assert all(isinstance(i, L.Symbol) for i in input)
+        assert all(isinstance(o, L.Symbol) for o in output)
         return L.Section(name, code, declaration, input, output, annotations)
 
     def _define_coordinate_dofs_lincomb(

ffcx/codegeneration/expression_generator.py

@@ -217,7 +217,7 @@ def generate_block_parts(self, blockmap, blockdata):
         components = ufl.product(self.ir.expression.shape)
 
         num_points = self.quadrature_rule[1].points.shape[0]
-        A_shape = [num_points, components] + self.ir.expression.tensor_shape
+        A_shape = [num_points, components] + list(self.ir.expression.tensor_shape)
         A = self.backend.symbols.element_tensor
         iq = self.backend.symbols.quadrature_loop_index
 

ffcx/codegeneration/integral_generator.py

@@ -13,6 +13,7 @@
 from typing import Any
 
 import basix
+import numpy as np
 import ufl
 
 import ffcx.codegeneration.lnodes as L
@@ -165,6 +166,10 @@ def generate(self, domain: basix.CellType):
         all_preparts = []
         all_quadparts = []
 
+        # Generate packing of proxy coefficients into contiguous arrays,
+        # will will be part of pre-computations.
+        all_preparts += self.generate_proxy_coefficient_packing()
+
         # Pre-definitions are collected across all quadrature loops to
         # improve re-use and avoid name clashes
         for cell, rule in self.ir.expression.integrand.keys():
@@ -313,6 +318,145 @@ def generate_piecewise_partition(self, quadrature_rule, domain: basix.CellType):
         arraysymbol = L.Symbol(f"sp_{quadrature_rule.id()}", dtype=L.DataType.SCALAR)
         return self.generate_partition(arraysymbol, F, "piecewise", None, None)
 
+    def generate_proxy_coefficient_packing(self):
+        """Generate packing of proxy coefficients into contiguous arrays."""
+        definitions = []
+        intermediates = []
+
+        proxy_coeff_offset = np.zeros(len(self.ir.proxy_coefficient_sizes) + 1, dtype=int)
+        proxy_coeff_offset[1:] = np.cumsum(self.ir.proxy_coefficient_sizes)
+        pw = L.Symbol("pw", dtype=L.DataType.SCALAR)
+        pw_array = L.ArrayDecl(pw, sizes=int(proxy_coeff_offset[-1]))
+
+        for i, (proxy_coeff, expr_name) in enumerate(self.ir.sub_expressions):
+            declarations = []
+
+            # Get active coefficients
+            active_coefficient_offsets = self.ir.proxy_coefficient_offsets[i : i + 2]
+            active_coefficients = self.ir.coefficients_in_proxy[
+                active_coefficient_offsets[0] : active_coefficient_offsets[1]
+            ]
+            sub_coefficient_sizes = [
+                active_coefficient.ufl_element().dim for active_coefficient in active_coefficients
+            ]
+
+            sub_coefficient_offsets = [
+                self.ir.expression.coefficient_offsets[coeff] for coeff in active_coefficients
+            ]
+            sub_coeff_pos = np.zeros(len(active_coefficients) + 1, dtype=np.int32)
+            sub_coeff_pos[1:] = np.cumsum(sub_coefficient_sizes)
+            # Declare array that holds subset of coefficients
+            sub_coeff = L.Symbol(f"sub_coeff_{i}", dtype=L.DataType.SCALAR)
+            sub_coeff_array = L.ArrayDecl(sub_coeff, sizes=int(np.sum(sub_coefficient_sizes)))
+            declarations.append(sub_coeff_array)
+
+            pi = L.Symbol("pi", dtype=L.DataType.INT)
+            # Pack coefficiets into contiguous array for expression evaluation
+            coeff_loops = []
+            for j in range(len(active_coefficients)):
+                coeff_loops.append(
+                    L.ForRange(
+                        pi,
+                        0,
+                        sub_coefficient_sizes[j],
+                        [
+                            L.Assign(
+                                sub_coeff[sub_coeff_pos[j] + pi],
+                                self.backend.symbols.coefficients[sub_coefficient_offsets[j] + pi],
+                            )
+                        ],
+                    )
+                )
+
+            pz_at_itg_points = L.Symbol(
+                f"proxy_coefficient_at_itg_points_{i}", dtype=L.DataType.SCALAR
+            )
+
+            # Initialize proxy coefficient array to zero
+            proxy_size = int(np.prod(self.ir.proxy_pack_shape[i]))
+            proxy_coefficient = L.ArrayDecl(pz_at_itg_points, sizes=proxy_size)
+            proxy_initialize = L.ForRange(pi, 0, proxy_size, [L.Assign(pz_at_itg_points[pi], 0.0)])
+            declarations.append(proxy_coefficient)
+
+            # NOTE: Need to do something similar for constants, currently we just pass them in
+            custom_data = L.Symbol("custom_data", dtype=L.DataType.SCALAR)
+            func_call = L.CallOp(
+                expr_name + ".tabulate_tensor",
+                (
+                    pz_at_itg_points,
+                    sub_coeff,
+                    self.backend.symbols.constants,
+                    self.backend.symbols.coordinate_dofs,
+                    self.backend.symbols.entity_local_index,
+                    self.backend.symbols.quadrature_permutation,
+                    custom_data,
+                ),
+            )
+            decl = L.Statement(func_call)
+            # Compute matvec between tabulated expression and interpolation matrix
+            identity_assign = False
+            if isinstance(proxy_coeff.operator, ufl.Interpolate):
+                be = proxy_coeff.ufl_element().basix_element
+                identity_assign = be.interpolation_is_identity
+                if not identity_assign:
+                    im = be.interpolation_matrix
+                    vs = int(np.prod(be.value_shape))
+            else:
+                raise NotImplementedError(
+                    "Only proxy coefficients for Interpolate supported at the moment"
+                )
+
+            num_dofs = self.ir.proxy_coefficient_sizes[i]
+            assign_start = proxy_coeff_offset[i]
+            if identity_assign:
+                inner_assign_loop = L.Assign(pw[assign_start + pi], pz_at_itg_points[pi])
+            else:
+                assert im.shape[0] == num_dofs
+                # Expression data is ordered xyzxyzxyz,
+                # Interpolation matrix is ordered xxxyyyzzz
+                if vs > 1:
+                    num_points = im.shape[1] // vs
+                    im_reshaped = im.reshape((im.shape[0], vs, num_points))
+                    im_transposed = im_reshaped.transpose((0, 2, 1))
+                    im = im_transposed.reshape((im.shape[0], -1))
+                im_table = self.declare_table(f"proxy_im_{i}", im)[0]
+                declarations.append(im_table)
+                num_quadrature_points = im.shape[1]
+                pj = L.Symbol("pj", dtype=L.DataType.INT)
+
+                inner_assign_loop = L.ForRange(
+                    pj,
+                    0,
+                    num_quadrature_points,
+                    [
+                        L.AssignAdd(
+                            pw[assign_start + pi], pz_at_itg_points[pj] * im_table.symbol[pi][pj]
+                        )
+                    ],
+                )
+            init_pw = L.Assign(pw[assign_start + pi], 0)
+            assign_loop = L.ForRange(pi, 0, num_dofs, [init_pw, inner_assign_loop])
+            intermediates += [
+                L.Section(
+                    f"Packing {i}th proxy coefficient",
+                    statements=[coeff_loops, proxy_initialize, decl],
+                    declarations=declarations,
+                    input=[],
+                    output=[],
+                )
+            ]
+            intermediates += [assign_loop]
+            intermediates = [
+                L.Section(
+                    name="Compute Proxy Coefficient",
+                    statements=intermediates,
+                    declarations=[pw_array],
+                    input=[],
+                    output=[],
+                )
+            ]
+        return definitions, intermediates
+
     def generate_varying_partition(self, quadrature_rule, domain: basix.CellType):
         """Generate a varying partition."""
         # Get annotated graph of factorisation

ffcx/codegeneration/jit.py

@@ -154,10 +154,10 @@ def compile_forms(
     options: dict = {},
     cache_dir: Path | None = None,
     timeout: int = 10,
-    cffi_extra_compile_args: list[str] = [],
+    cffi_extra_compile_args: list[str] | None = None,
     cffi_verbose: bool = False,
     cffi_debug: bool = False,
-    cffi_libraries: list[str] = [],
+    cffi_libraries: list[str] | None = None,
     visualise: bool = False,
 ):
     """Compile a list of UFL forms into UFCx Python objects.
@@ -174,6 +174,8 @@ def compile_forms(
         visualise: Toggle visualisation
     """
     p = ffcx.options.get_options(options)
+    cffi_extra_compile_args = cffi_extra_compile_args or []
+    cffi_libraries = cffi_libraries or []
 
     # If requested, replace bi-linear forms by their diagonal part
     if p["part"] == "diagonal":
@@ -248,11 +250,11 @@ def compile_forms(
 
 
 def compile_expressions(
-    expressions: list[tuple[ufl.Expr, npt.NDArray[np.floating]]],  # type: ignore
-    options: dict = {},
+    expressions: list[tuple[ufl.core.expr.Expr, npt.NDArray[np.floating]]],
+    options: dict | None = None,
     cache_dir: Path | None = None,
     timeout: int = 10,
-    cffi_extra_compile_args: list[str] = [],
+    cffi_extra_compile_args: list[str] | None = None,
     cffi_verbose: bool = False,
     cffi_debug: bool = False,
     cffi_libraries: list[str] = [],