diff --git a/ffcx/analysis.py b/ffcx/analysis.py index 32e847bec..69959cd32 100644 --- a/ffcx/analysis.py +++ b/ffcx/analysis.py @@ -1,4 +1,4 @@ -# Copyright (C) 2007-2020 Anders Logg, Martin Alnaes, Kristian B. Oelgaard, +# Copyright (C) 2007-2026 Anders Logg, Martin Alnaes, Kristian B. Oelgaard, # Michal Habera and others # # This file is part of FFCx. (https://www.fenicsproject.org) @@ -15,6 +15,7 @@ import logging import typing +from functools import singledispatchmethod if typing.TYPE_CHECKING: from ufl.algorithms.formdata import FormData @@ -23,10 +24,55 @@ import numpy as np import numpy.typing as npt import ufl.algorithms +from ufl.algorithms.apply_derivatives import apply_coordinate_derivatives, apply_derivatives +from ufl.algorithms.apply_function_pullbacks import apply_function_pullbacks +from ufl.algorithms.apply_geometry_lowering import apply_geometry_lowering +from ufl.algorithms.apply_integral_scaling import apply_integral_scaling +from ufl.algorithms.compute_form_data import attach_estimated_degrees, preprocess_form + +# See TODOs at the call sites of these below: +from ufl.algorithms.domain_analysis import ( + build_integral_data, + group_form_integrals, +) +from ufl.algorithms.formdata import FormData +from ufl.algorithms.remove_complex_nodes import remove_complex_nodes +from ufl.algorithms.remove_component_tensors import remove_component_tensors +from ufl.corealg.dag_traverser import DAGTraverser logger = logging.getLogger("ffcx") +def apply_push_forward( + expr: ufl.core.expr.Expr, domain: ufl.Mesh, pullback: ufl.pullback.AbstractPullback +) -> ufl.core.expr.Expr: + """Apply push forward to an expression.""" + J = ufl.Jacobian(domain) + J_T = J.T + detJ = ufl.JacobianDeterminant(domain) + invJ = ufl.JacobianInverse(domain) + invJ_T = invJ.T + pushed_forward_expr: ufl.core.expr.Expr + match pullback: + case ufl.pullback.L2Piola(): + pushed_forward_expr = detJ * expr + case ufl.pullback.CovariantPiola(): + pushed_forward_expr = ufl.dot(J_T, expr) + case ufl.pullback.ContravariantPiola(): + pushed_forward_expr = detJ * ufl.dot(invJ, expr) + case ufl.pullback.DoubleCovariantPiola(): + pushed_forward_expr = ufl.dot(J_T, ufl.dot(expr, J)) + case ufl.pullback.DoubleContravariantPiola(): + pushed_forward_expr = detJ * ufl.dot(invJ, ufl.dot(expr, invJ_T)) + case ufl.pullback.CovariantContravariantPiola(): + pushed_forward_expr = detJ * ufl.dot(J_T, ufl.dot(expr, invJ_T)) + case ufl.pullback.IdentityPullback(): + pushed_forward_expr = expr + case _: + raise NotImplementedError(f"Pullback {pullback} not supported.") + return pushed_forward_expr + + class UFLData(typing.NamedTuple): """UFL data.""" @@ -93,6 +139,59 @@ def analyze_ufl_objects( elements += data.unique_sub_elements coordinate_elements += data.coordinate_elements + # Loop through forms to extract interpolate operands + new_coefficients = [] + for data in form_data: + current_coeffs = data.reduced_coefficients.copy() + for coeff in current_coeffs: + if isinstance(coeff, ProxyCoefficient): + # Expose expression used for interpolation to generated code + original_expression = coeff.operand + element = coeff.ufl_function_space().ufl_element() + # Push expression forward to physical cell + domain = coeff.ufl_function_space().ufl_domain() + pushed_forward_expression = apply_push_forward( + original_expression, domain, element.pullback + ) + elements += ufl.algorithms.extract_elements(pushed_forward_expression) + processed_expression = _analyze_expression(pushed_forward_expression, scalar_type) + points = element.basix_element.points + processed_expressions += [(processed_expression, points, original_expression)] + + # Append coefficents in the processed expression + # to the form data reduced coefficients. + new_coefficients.extend( + [ + coeff + for coeff in ufl.algorithms.extract_coefficients(processed_expression) + if coeff not in data.reduced_coefficients + ] + ) + data._reduced_coefficients.extend(new_coefficients) + + # Update form data for new set of reduced coefficients + # Sort the reduced coefficients of the form. + data._reduced_coefficients = sorted(data._reduced_coefficients, key=lambda x: x.count()) + # NOTE: Could have been simpler if we had extracted the elements from + # reduced_coefficients rather than going through coefficient_elements. + new_coeff_elements = tuple(coeff.ufl_element() for coeff in data.reduced_coefficients) + data._coefficient_elements = new_coeff_elements + # Enable all coefficients that are either in the original integral coefficients or + # in the new coefficients generated from the interpolate expressions. + for itg_data in data.integral_data: + assert itg_data.integral_coefficients is not None + itg_data.enabled_coefficients = [ + bool(coeff in itg_data.integral_coefficients) or bool(coeff in new_coefficients) + for coeff in data.reduced_coefficients + ] + + # Update original coefficient position in form + data._original_coefficient_positions = [ + i + for i, c in enumerate(data.original_form.coefficients()) + if c in data.reduced_coefficients + ] + for original_expression, points in expressions: elements += ufl.algorithms.extract_elements(original_expression) processed_expression = _analyze_expression(original_expression, scalar_type) @@ -172,17 +271,21 @@ def _analyze_form(form: ufl.Form, scalar_type: npt.DTypeLike) -> FormData: complex_mode = np.issubdtype(scalar_type, np.complexfloating) # Compute form metadata - form_data: FormData = ufl.algorithms.compute_form_data( + form_data: FormData = compute_form_data( form, do_apply_function_pullbacks=True, do_apply_integral_scaling=True, do_apply_geometry_lowering=True, - preserve_geometry_types=(ufl.classes.Jacobian,), + preserve_geometry_types=(ufl.geometry.Jacobian,), # type: ignore do_apply_restrictions=True, do_append_everywhere_integrals=False, # do not add dx integrals to dx(i) in UFL complex_mode=complex_mode, ) + # Store original form, prior to replacement of interpolate, + # to ensure we get the correct signature and coefficient numbering + form_data._original_form = form + # Determine unique quadrature degree and quadrature scheme # per each integral data for id, integral_data in enumerate(form_data.integral_data): @@ -259,3 +362,229 @@ def _has_custom_integrals( return any(_has_custom_integrals(itg) for itg in o) else: raise NotImplementedError + + +class ProxyCoefficient(ufl.Coefficient): + """Proxy coefficient to replace operands that require custom treatement.""" + + _operator: ufl.core.expr.Expr + + def __init__(self, V: ufl.FunctionSpace, operator: ufl.core.expr.Expr): + """Initialise.""" + self._operator = operator + super().__init__(V) + + @property + def operand(self) -> ufl.core.expr.Expr: + """The operand that this proxy coefficient is replacing.""" + return self._operator.ufl_operands[0] + + @property + def operator(self) -> ufl.core.expr.Expr: + """The kind of the operand that this proxy coefficient is replacing.""" + return self._operator + + +class IntermediateCoefficientReplacer(DAGTraverser): + """Replace operands requiring intermediate coefficients with intermediate objects.""" + + def __init__( + self, + compress: bool | None = True, + visited_cache: dict[tuple, ufl.core.expr.Expr] | None = None, + result_cache: dict[ufl.core.expr.Expr, ufl.core.expr.Expr] | None = None, + ) -> None: + """Initialise. + + Args: + compress: If True, ``result_cache`` will be used. + visited_cache: cache of intermediate results; + expr -> r = self.process(expr, ...). + result_cache: cache of result objects for memory reuse, r -> r. + + """ + super().__init__(compress=compress, visited_cache=visited_cache, result_cache=result_cache) + + @singledispatchmethod + def process( + self, + o: ufl.core.expr.Expr, + reference_value: bool | None = False, + reference_grad: int | None = 0, + restricted: str | None = None, + ) -> ufl.core.expr.Expr: + """Replace averaged arguments with intermediate space. + + Args: + o: `ufl.core.expr.Expr` to be processed. + reference_value: Whether `ReferenceValue` has been applied or not. + reference_grad: Number of `ReferenceGrad`s that have been applied. + restricted: '+', '-', or None. + """ + return super().process(o) + + @process.register(ufl.Interpolate) + def _( + self, + o: ufl.Interpolate, + reference_value: bool | None = False, + reference_grad: int | None = 0, + restricted: str | None = None, + ) -> ufl.core.expr.Expr: + """Handle Interpolate.""" + ops = o.ufl_operands + assert len(ops) == 1, "Expected single operator in interpolation" + return ProxyCoefficient(o.ufl_function_space(), o) + + @process.register(ufl.core.expr.Expr) + def _( + self, + o: ufl.Argument, + reference_value: bool | None = False, + reference_grad: int | None = 0, + restricted: str | None = None, + ) -> ufl.core.expr.Expr: + """Handle anything else in UFL.""" + return self.reuse_if_untouched( + o, + reference_value=reference_value, + reference_grad=reference_grad, + restricted=restricted, + ) + + +def replace_ufl_operands(form: ufl.Form) -> ufl.Form: + """Parse UFL form and and add placeholder coefficients. + + Args: + form: UFL form + + Return: + The modified form with operands replaced. + """ + rule = IntermediateCoefficientReplacer() + return ufl.algorithms.map_integrands.map_integrands(rule, form) # type: ignore + + +def compute_form_data( + form: ufl.Form, + do_apply_function_pullbacks: bool = False, + do_apply_integral_scaling: bool = False, + do_apply_geometry_lowering: bool = False, + preserve_geometry_types: tuple[ufl.geometry.GeometricQuantity, ...] = (), + do_apply_default_restrictions: bool = True, + do_apply_restrictions: bool = True, + do_estimate_degrees: bool = True, + do_append_everywhere_integrals: bool = True, + do_replace_functions: bool = False, + coefficients_to_split: tuple[ufl.Coefficient, ...] | None = None, + complex_mode: bool = False, + do_remove_component_tensors: bool = False, +) -> FormData: + """Compute form data. + + Args: + form: The form to compute form data for. + do_apply_function_pullbacks: Apply pull-back to reference cell + for coefficients, including Piola and symmetry transforms + if required. + do_apply_integral_scaling: Apply scaling of moving the integral + from physical to reference frame. + do_apply_geometry_lowering: Lower the representation of geometrical + quantities to a smaller subset of quantities + preserve_geometry_types: Set of quantities not to lower, and keep + at its present stage for the form-compiler. + do_apply_default_restrictions: Apply default restrictions, defined in + {py:mod}`ufl.algorithms.apply_restrictions` to integrals if no + restriction has been set. + do_apply_restrictions: Apply restrictions towards terminal nodes. + do_replace_functions: Replace functions with with its cannonically numbered + function or thos provided in coefficients_to_split. + coefficients_to_split: Sequence of coefficients to split over a MeshSequence. + do_estimate_degrees: Estimate polynomial degree of integrands. + do_append_everywhere_integrals: If True append every `dx` integral to each `dx(i)` + integral defined in the form. + do_remove_component_tensors: Remove component-tensor if true. + complex_mode: If false remove complex nodes from the form. + """ + # --- Store untouched form for reference. + # The user of FormData may get original arguments, + # original coefficients, and form signature from this object. + # But be aware that the set of original coefficients are not + # the same as the ones used in the final UFC form. + # See 'reduced_coefficients' below. + original_form = form + + # --- Pass form integrands through some symbolic manipulation + form = replace_ufl_operands(form) + + form = preprocess_form(form, complex_mode) + + # --- Group form integrals + # TODO: Refactor this, it's rather opaque what this does + # TODO: Is self.original_form.ufl_domains() right here? + # It will matter when we start including 'num_domains' in ufc form. + form = group_form_integrals( + form, + original_form.ufl_domains(), + do_append_everywhere_integrals=do_append_everywhere_integrals, + ) + + # Estimate polynomial degree of integrands now, before applying + # any pullbacks and geometric lowering. Otherwise quad degrees + # blow up horrifically. + if do_estimate_degrees: + form = attach_estimated_degrees(form) + + if do_apply_function_pullbacks: + # Rewrite coefficients and arguments in terms of their + # reference cell values with Piola transforms and symmetry + # transforms injected where needed. + # Decision: Not supporting grad(dolfin.Expression) without a + # Domain. Current dolfin works if Expression has a + # cell but this should be changed to a mesh. + form = apply_function_pullbacks(form) + + # Scale integrals to reference cell frames + if do_apply_integral_scaling: + form = apply_integral_scaling(form) + + # Lower abstractions for geometric quantities into a smaller set + # of quantities, allowing the form compiler to deal with a smaller + # set of types and treating geometric quantities like any other + # expressions w.r.t. loop-invariant code motion etc. + if do_apply_geometry_lowering: + form = apply_geometry_lowering(form, preserve_geometry_types) + + # Apply differentiation again, because the algorithms above can + # generate new derivatives or rewrite expressions inside + # derivatives + if do_apply_function_pullbacks or do_apply_geometry_lowering: + form = apply_derivatives(form) + + # Neverending story: apply_derivatives introduces new Jinvs, + # which needs more geometry lowering + if do_apply_geometry_lowering: + form = apply_geometry_lowering(form, preserve_geometry_types) + # Lower derivatives that may have appeared + form = apply_derivatives(form) + + form = apply_coordinate_derivatives(form) + + # If in real mode, remove any complex nodes introduced during form processing. + if not complex_mode: + form = remove_complex_nodes(form) + + # Remove component tensors + if do_remove_component_tensors: + form = remove_component_tensors(form) + integral_data = build_integral_data(form.integrals()) + return FormData( + original_form, + integral_data, + do_apply_default_restrictions=do_apply_default_restrictions, + do_apply_restrictions=do_apply_restrictions, + do_replace_functions=do_replace_functions, + coefficients_to_split=coefficients_to_split, + complex_mode=complex_mode, + ) diff --git a/ffcx/codegeneration/C/formatter.py b/ffcx/codegeneration/C/formatter.py index b5845366a..41686309e 100644 --- a/ffcx/codegeneration/C/formatter.py +++ b/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.""" diff --git a/ffcx/codegeneration/access.py b/ffcx/codegeneration/access.py index 2989fafb4..0da5a4cb8 100644 --- a/ffcx/codegeneration/access.py +++ b/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, diff --git a/ffcx/codegeneration/backend.py b/ffcx/codegeneration/backend.py index ef6963987..63b1f3bc3 100644 --- a/ffcx/codegeneration/backend.py +++ b/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 diff --git a/ffcx/codegeneration/codegeneration.py b/ffcx/codegeneration/codegeneration.py index 25e221b01..f63d0506d 100644 --- a/ffcx/codegeneration/codegeneration.py +++ b/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]] diff --git a/ffcx/codegeneration/definitions.py b/ffcx/codegeneration/definitions.py index bfe517030..14e77ac62 100644 --- a/ffcx/codegeneration/definitions.py +++ b/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( diff --git a/ffcx/codegeneration/expression_generator.py b/ffcx/codegeneration/expression_generator.py index a4aca9af8..8710561a3 100644 --- a/ffcx/codegeneration/expression_generator.py +++ b/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 diff --git a/ffcx/codegeneration/integral_generator.py b/ffcx/codegeneration/integral_generator.py index a217b50fe..1c36130b5 100644 --- a/ffcx/codegeneration/integral_generator.py +++ b/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 diff --git a/ffcx/codegeneration/jit.py b/ffcx/codegeneration/jit.py index 0b8904c2c..2b78f6a27 100644 --- a/ffcx/codegeneration/jit.py +++ b/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] = [], diff --git a/ffcx/codegeneration/lnodes.py b/ffcx/codegeneration/lnodes.py index debd4441e..ea766075b 100644 --- a/ffcx/codegeneration/lnodes.py +++ b/ffcx/codegeneration/lnodes.py @@ -482,6 +482,17 @@ def __init__(self, lhs, rhs): self.dtype = merge_dtypes([self.lhs.dtype, self.rhs.dtype]) +class CallOp(LExprOperator): + """A function call operator.""" + + op = "" + + def __init__(self, function, args): + """Initialise.""" + self.func = function + self.args = [as_lexpr(arg) for arg in args] + + class NaryOp(LExprOperator): """Base class for special n-ary operators.""" diff --git a/ffcx/codegeneration/symbols.py b/ffcx/codegeneration/symbols.py index 2e4830fc1..957acff81 100644 --- a/ffcx/codegeneration/symbols.py +++ b/ffcx/codegeneration/symbols.py @@ -65,13 +65,23 @@ def format_mt_name(basename, mt): class FFCXBackendSymbols: """FFCx specific symbol definitions. Provides non-ufl symbols.""" - def __init__(self, coefficient_numbering, coefficient_offsets, original_constant_offsets): + def __init__( + self, + coefficient_numbering, + coefficient_offsets, + original_constant_offsets, + proxy_coefficient_numbering, + proxy_coefficient_offsets, + ): """Initialise.""" self.coefficient_numbering = coefficient_numbering self.coefficient_offsets = coefficient_offsets self.original_constant_offsets = original_constant_offsets + self.proxy_coefficient_numbering = proxy_coefficient_numbering + self.proxy_coefficient_offsets = proxy_coefficient_offsets + # Keep tabs on tables, so the symbols can be reused self.quadrature_weight_tables = {} self.element_tables = {} @@ -82,6 +92,7 @@ def __init__(self, coefficient_numbering, coefficient_offsets, original_constant # Symbols for the tabulate_tensor function arguments self.element_tensor = L.Symbol("A", dtype=L.DataType.SCALAR) self.coefficients = L.Symbol("w", dtype=L.DataType.SCALAR) + self.proxy_coefficients = L.Symbol("pw", dtype=L.DataType.SCALAR) self.constants = L.Symbol("c", dtype=L.DataType.SCALAR) self.coordinate_dofs = L.Symbol("coordinate_dofs", dtype=L.DataType.REAL) self.entity_local_index = L.Symbol("entity_local_index", dtype=L.DataType.INT) @@ -173,6 +184,28 @@ def coefficient_value(self, mt): c = self.coefficient_numbering[mt.terminal] return L.Symbol(format_mt_name(f"w{c:d}", mt), dtype=L.DataType.SCALAR) + def proxy_coefficient_dof_access(self, coefficient, dof_index): + """Coefficient DOF access.""" + offset = self.proxy_coefficient_offsets[coefficient] + z = self.proxy_coefficients + return z[offset + dof_index] + + def proxy_coefficient_value(self, mt): + """Symbol for variable holding value or derivative component of proxy coefficient.""" + c = self.proxy_coefficient_numbering[mt.terminal] + return L.Symbol(format_mt_name(f"pw{c:d}", mt), dtype=L.DataType.SCALAR) + + def proxy_coefficient_dof_access_blocked( + self, coefficient: ufl.Coefficient, index, block_size, dof_offset + ): + """Blocked proxy coefficient DOF access.""" + coeff_offset = self.coefficient_offsets[coefficient] + z = self.proxy_coefficients + _z = L.Symbol(f"_pw_{coeff_offset}_{dof_offset}", dtype=L.DataType.SCALAR) + unit_stride_access = _z[index] + original_access = z[coeff_offset + index * block_size + dof_offset] + return unit_stride_access, original_access + def constant_index_access(self, constant, index): """Constant index access.""" offset = self.original_constant_offsets[constant] diff --git a/ffcx/ir/analysis/graph.py b/ffcx/ir/analysis/graph.py index 72bab5a99..7e6086d96 100644 --- a/ffcx/ir/analysis/graph.py +++ b/ffcx/ir/analysis/graph.py @@ -1,4 +1,4 @@ -# Copyright (C) 2011-2017 Martin Sandve Alnæs +# Copyright (C) 2011-2026 Martin Sandve Alnæs # # This file is part of FFCx.(https://www.fenicsproject.org) # @@ -27,9 +27,9 @@ class ExpressionGraph: def __init__(self): """Initialise.""" # Data structures for directed multi-edge graph - self.nodes = {} - self.out_edges = {} - self.in_edges = {} + self._nodes = {} + self._out_edges = {} + self._in_edges = {} self.e2i = {} def number_of_nodes(self): @@ -38,17 +38,32 @@ def number_of_nodes(self): def add_node(self, key, **kwargs): """Add a node with optional properties.""" - self.nodes[key] = kwargs - self.out_edges[key] = [] - self.in_edges[key] = [] + self._nodes[key] = kwargs + self._out_edges[key] = [] + self._in_edges[key] = [] def add_edge(self, node1, node2): """Add a directed edge from node1 to node2.""" if node1 not in self.nodes or node2 not in self.nodes: raise KeyError("Adding edge to unknown node") - self.out_edges[node1] += [node2] - self.in_edges[node2] += [node1] + self._out_edges[node1] += [node2] + self._in_edges[node2] += [node1] + + @property + def nodes(self): + """Get nodes.""" + return self._nodes + + @property + def out_edges(self): + """Get out edges.""" + return self._out_edges + + @property + def in_edges(self): + """Get in edges.""" + return self._in_edges def build_graph_vertices(expressions, skip_terminal_modifiers=False) -> ExpressionGraph: diff --git a/ffcx/ir/analysis/modified_terminals.py b/ffcx/ir/analysis/modified_terminals.py index 528cfd748..8e6c0681c 100644 --- a/ffcx/ir/analysis/modified_terminals.py +++ b/ffcx/ir/analysis/modified_terminals.py @@ -276,6 +276,7 @@ def analyse_modified_terminal(expr): # Get the shape of the core terminal or its reference value, this is # the shape that component refers to if isinstance(t, FormArgument): + assert hasattr(t, "ufl_function_space") element = t.ufl_function_space().ufl_element() if reference_value: # Ignoring symmetry, assuming already applied in conversion diff --git a/ffcx/ir/elementtables.py b/ffcx/ir/elementtables.py index e2108a859..e12ba7bad 100644 --- a/ffcx/ir/elementtables.py +++ b/ffcx/ir/elementtables.py @@ -255,7 +255,8 @@ def get_modified_terminal_element(mt) -> ModifiedTerminalElement | None: raise RuntimeError("Global derivatives of reference values not defined.") elif ld and not mt.reference_value: raise RuntimeError("Local derivatives of global values not defined.") - element = mt.terminal.ufl_function_space().ufl_element() # type: ignore + assert hasattr(mt.terminal, "ufl_function_space") + element = mt.terminal.ufl_function_space().ufl_element() fc = mt.flat_component elif isinstance(mt.terminal, ufl.classes.SpatialCoordinate): if mt.reference_value: @@ -350,7 +351,7 @@ def permute_quadrature_quadrilateral(points, reflections=0, rotations=0): def build_optimized_tables( quadrature_rule: QuadratureRule, - cell: ufl.Cell, + cell: ufl.Cell | None, integral_type: typing.Literal["interior_facet", "exterior_facet", "ridge", "cell", "vertex"], entity_type: entity_types, modified_terminals: typing.Iterable[ModifiedTerminal], @@ -425,6 +426,7 @@ def build_optimized_tables( # It should be possible to reuse the cached tables by name, but # the dofmap offset may differ due to restriction. + assert isinstance(cell, ufl.Cell), "Cell must be a ufl.Cell object." tdim = cell.topological_dimension codim = tdim - element.cell.topological_dimension assert codim >= 0 diff --git a/ffcx/ir/integral.py b/ffcx/ir/integral.py index 877573647..4d02a88cd 100644 --- a/ffcx/ir/integral.py +++ b/ffcx/ir/integral.py @@ -19,6 +19,7 @@ from ufl.checks import is_cellwise_constant from ufl.classes import QuadratureWeight +from ffcx.analysis import ProxyCoefficient from ffcx.definitions import entity_types, supported_integral_types from ffcx.ir.analysis.factorization import compute_argument_factorization from ffcx.ir.analysis.graph import ExpressionGraph, build_scalar_graph @@ -103,9 +104,9 @@ class IntermediateIntegralIR(typing.TypedDict): """Intermediate IR for integrals.""" needs_facet_permutations: bool - unique_tables: dict[basix.CellType, dict[str, npt.NDArray[np.float64]]] - unique_table_types: dict[basix.CellType, dict[str, _table_types]] - integrand: dict[tuple[basix.CellType, QuadratureRule], IntermediateIntegrandIR] + unique_tables: dict[basix.CellType | str, dict[str, npt.NDArray[np.float64]]] + unique_table_types: dict[basix.CellType | str, dict[str, _table_types]] + integrand: dict[tuple[basix.CellType | str, QuadratureRule], IntermediateIntegrandIR] class CommonExpressionIR(typing.NamedTuple): @@ -121,6 +122,14 @@ class CommonExpressionIR(typing.NamedTuple): coefficient_numbering: dict[ufl.Coefficient, int] # Each coefficient is accessed from a given offset in the input array coefficient_offsets: dict[ufl.Coefficient, int] + + proxy_coefficient_numbering: dict[ + ProxyCoefficient, int + ] # For coefficients that are accessed via a proxy + proxy_coefficient_offsets: dict[ + ProxyCoefficient, int + ] # For coefficients that are accessed via a proxy + original_constant_offsets: dict[ufl.Constant, int] unique_tables: dict[basix.CellType, npt.NDArray[np.float64]] unique_table_types: dict[basix.CellType, dict[str, str]] @@ -136,7 +145,7 @@ def _compute_integral_ir( expression: ufl.core.expr.Expr, existing_tables: dict[str, npt.NDArray[np.float64]], quadrature_rule: QuadratureRule, - cell: ufl.Cell, + cell: ufl.Cell | None, integral_type: supported_integral_types, entity_type: entity_types, argument_shape: tuple[int, ...], @@ -172,6 +181,7 @@ def _compute_integral_ir( # Check if we have a mixed-dimensional integral is_mixed_dim = False for domain in ufl.domain.extract_domains(expression): + assert isinstance(cell, ufl.Cell), "Cell must be a ufl.Cell object." if domain.topological_dimension != cell.topological_dimension: is_mixed_dim = True @@ -385,10 +395,10 @@ def _compute_integral_ir( def compute_integral_ir( - cell: ufl.Cell, + cell: ufl.Cell | None, integral_type: supported_integral_types, entity_type: entity_types, - integrands: dict[basix.CellType, dict[QuadratureRule, ufl.core.expr.Expr]], + integrands: dict[basix.CellType | str, dict[QuadratureRule, ufl.core.expr.Expr]], argument_shape: tuple[int, ...], p: dict, visualise: bool, @@ -407,9 +417,9 @@ def compute_integral_ir( """ # Data that will populate the intermediate representation. needs_facet_permutations: bool = False - unique_tables: dict[basix.CellType, dict[str, npt.NDArray[np.float64]]] = {} - unique_table_types: dict[basix.CellType, dict[str, _table_types]] = {} - integrand_map: dict[tuple[basix.CellType, QuadratureRule], IntermediateIntegrandIR] = {} + unique_tables: dict[basix.CellType | str, dict[str, npt.NDArray[np.float64]]] = {} + unique_table_types: dict[basix.CellType | str, dict[str, _table_types]] = {} + integrand_map: dict[tuple[basix.CellType | str, QuadratureRule], IntermediateIntegrandIR] = {} for integral_domain, integrands_on_domain in integrands.items(): unique_tables[integral_domain] = {} unique_table_types[integral_domain] = {} @@ -502,6 +512,7 @@ def analyse_dependencies(F, mt_unique_table_reference): raise RuntimeError(f"Invalid ttype {ttype}.") elif not is_cellwise_constant(v["expression"]): + breakpoint() raise RuntimeError("Error " + str(tr)) # Keeping this check to be on the safe side, # not sure which cases this will cover (if any) diff --git a/ffcx/ir/representation.py b/ffcx/ir/representation.py index 4f8ac5221..9e63d7770 100644 --- a/ffcx/ir/representation.py +++ b/ffcx/ir/representation.py @@ -35,7 +35,7 @@ from ufl.sorting import sorted_expr_sum from ffcx import naming -from ffcx.analysis import UFLData +from ffcx.analysis import ProxyCoefficient, UFLData from ffcx.definitions import entity_types, supported_integral_types from ffcx.ir.integral import CommonExpressionIR, TensorPart, compute_integral_ir from ffcx.ir.representationutils import ( @@ -106,6 +106,11 @@ class IntegralIR(typing.NamedTuple): rank: int enabled_coefficients: list[bool] part: TensorPart + sub_expressions: list[tuple[ProxyCoefficient, str]] + proxy_coefficient_sizes: list[int] + proxy_pack_shape: list[tuple[int, ...]] + coefficients_in_proxy: list[ufl.Coefficient] + proxy_coefficient_offsets: list[int] class ExpressionIR(typing.NamedTuple): @@ -257,12 +262,17 @@ def compute_ir( prefix, ) + expression_names = { + org_expr: naming.expression_name((org_expr, points), prefix) + for (expr, points, org_expr) in analysis.expressions + } irs = [ _compute_integral_ir( fd, i, analysis.element_numbers.keys(), integral_names, + expression_names, options, visualise, ) @@ -287,7 +297,6 @@ def compute_ir( ) for (i, fd) in enumerate(analysis.form_data) ] - ir_expressions = [ _compute_expression_ir( expr, @@ -300,7 +309,6 @@ def compute_ir( ) for i, expr in enumerate(analysis.expressions) ] - return DataIR( integrals=ir_integrals, forms=ir_forms, @@ -313,6 +321,7 @@ def _compute_integral_ir( form_index: int, unique_elements: typing.Iterable[basix.ufl._ElementBase], integral_names: dict[tuple[int, int], str], + expression_names: dict[ufl.core.expr.Expr, str], options, visualise, ) -> list[IntegralIR]: @@ -323,6 +332,8 @@ def _compute_integral_ir( form_index: Index of form in the sequence of forms. unique_elements: Set of unique elements in the form. integral_names: Map from `(form_index, integral_index)` to the name of the integral. + expression_names: Map from original expression to the name of the expression. + Used for sub-expressions that are coefficients in the integral (internal proxies). options: Options for the intermediate representation. 'part': If the full tensor or the diagonal of the tensor should be generated. Only valid for bi-linear forms. 'sum_factorization': If sum factorization should be used. Only has an effect on cell @@ -437,18 +448,32 @@ def _compute_integral_ir( # input coefficient data a coefficient should start. coefficient_numbering: dict[ufl.Coefficient, int] = {} coefficient_offsets: dict[ufl.Coefficient, int] = {} - _offset = 0 + proxy_coefficient_numbering: dict[ProxyCoefficient, int] = {} + _proxy_coefficient_offsets: dict[ProxyCoefficient, int] = {} + _offset_c = 0 + _offset_p = 0 width = 2 if integral_type in ("interior_facet") else 1 + first_proxy = True + start_proxy = len(form_data.reduced_coefficients) for i, (coeff, el) in enumerate( zip(form_data.reduced_coefficients, form_data.coefficient_elements) ): - coefficient_numbering[coeff] = i - coefficient_offsets[coeff] = _offset - _offset += width * element_dimensions[el] + if isinstance(coeff, ProxyCoefficient): + proxy_coefficient_numbering[coeff] = i + _proxy_coefficient_offsets[coeff] = _offset_p + _offset_p += width * element_dimensions[el] + start_proxy = i if first_proxy else start_proxy + first_proxy = False + else: + coefficient_numbering[coeff] = i + coefficient_offsets[coeff] = _offset_c + _offset_c += width * element_dimensions[el] # Add Coefficient numbering and offsets to IR expression_ir["coefficient_numbering"] = coefficient_numbering expression_ir["coefficient_offsets"] = coefficient_offsets + expression_ir["proxy_coefficient_numbering"] = proxy_coefficient_numbering + expression_ir["proxy_coefficient_offsets"] = _proxy_coefficient_offsets # Similar procedure with constants, but they are not removed # as they usually contain very little data. @@ -461,7 +486,7 @@ def _compute_integral_ir( expression_ir["original_constant_offsets"] = original_constant_offsets # Create map from number of quadrature points -> integrand - integrand_map: dict[basix.CellType, dict[QuadratureRule, Expr]] = { + integrand_map: dict[basix.CellType | str, dict[QuadratureRule, Expr]] = { cell_type: {rule: integral.integrand() for rule, integral in cell_integrals.items()} for cell_type, cell_integrals in sorted_integrals.items() } @@ -477,6 +502,44 @@ def _compute_integral_ir( visualise, ) + # Check if the coefficients of the integral is a proxy coefficient, and if it is enabled. + proxy_coefficients: list[ufl.Coefficient] = [] + if itg_data.integral_coefficients is not None and itg_data.enabled_coefficients is not None: + enabled_itg_coeffcients = [ + c + for (i, c) in enumerate(itg_data.integral_coefficients) + if itg_data.enabled_coefficients[i] + ] + proxy_coefficients = [ + coeff + for coeff in form_data.reduced_coefficients[start_proxy:] + if coeff in enabled_itg_coeffcients + ] + + # Num coefficients required for the form, including generating proxy coefficients + proxy_coefficient_offsets = [0] + ir["coefficients_in_proxy"] = [] + ir["sub_expressions"] = [] + for p_coeff in proxy_coefficients: + assert isinstance(p_coeff, ProxyCoefficient) + for coeff in ufl.algorithms.extract_coefficients(p_coeff.operand): + ir["coefficients_in_proxy"].append(coeff) + proxy_coefficient_offsets.append(len(ir["coefficients_in_proxy"])) + ir["sub_expressions"].append((p_coeff, expression_names[p_coeff.operand])) + ir["proxy_coefficient_offsets"] = proxy_coefficient_offsets + + ir["proxy_pack_shape"] = [] + ir["proxy_coefficient_sizes"] = [] + for p_coeff in proxy_coefficients: + value_size = int(np.prod(p_coeff.ufl_shape)) + ir["proxy_coefficient_sizes"].append(p_coeff.ufl_element().dim) + ir["proxy_pack_shape"].append( + ( + value_size, + p_coeff.ufl_function_space().ufl_element().basix_element.points.shape[0], + ) + ) + expression_ir.update(integral_ir) # Fetch name @@ -516,7 +579,7 @@ def _compute_form_ir( else: ir["rank"] = len(form_data.original_form.arguments()) - ir["num_coefficients"] = len(form_data.reduced_coefficients) + ir["num_coefficients"] = len(form_data.original_coefficient_positions) ir["coefficient_names"] = [ object_names.get(id(obj), f"w{j}") for j, obj in enumerate(form_data.reduced_coefficients) @@ -572,27 +635,27 @@ def _compute_form_ir( def _compute_expression_ir( - expr, - index, - prefix, - analysis, - options, - visualise, + analysis_expr: tuple[ufl.core.expr.Expr, np.ndarray, ufl.core.expr.Expr], + index: int, + prefix: str, + analysis: UFLData, + options: dict, + visualise: bool, object_names, ): """Compute intermediate representation of an Expression.""" logger.info(f"Computing IR for Expression {index}") # Compute representation - ir = {} - base_ir = {} - original_expr = (expr[2], expr[1]) + ir: dict[str, typing.Any] = {} + base_ir: dict[str, typing.Any] = {} + original_input_expr = (analysis_expr[2], analysis_expr[1]) - base_ir["name"] = naming.expression_name(original_expr, prefix) + base_ir["name"] = naming.expression_name(original_input_expr, prefix) - original_expr = expr[2] - points = expr[1] - expr = expr[0] + original_expr = analysis_expr[2] + points = analysis_expr[1] + expr = analysis_expr[0] expr_domain = max( ufl.domain.extract_domains(expr), default=None, key=lambda d: d.topological_dimension @@ -606,30 +669,13 @@ def _compute_expression_ir( # Extract dimensions for elements of arguments only arguments = ufl.algorithms.extract_arguments(expr) argument_elements = tuple(f.ufl_function_space().ufl_element() for f in arguments) - argument_dimensions = [element_dimensions[element] for element in argument_elements] + argument_dimensions = tuple(element_dimensions[element] for element in argument_elements) tensor_shape = argument_dimensions base_ir["tensor_shape"] = tensor_shape base_ir["shape"] = list(expr.ufl_shape) - coefficients = ufl.algorithms.extract_coefficients(expr) - coefficient_numbering = {} - for i, coeff in enumerate(coefficients): - coefficient_numbering[coeff] = i - - # Add coefficient numbering to IR - base_ir["coefficient_numbering"] = coefficient_numbering - - original_coefficient_positions = [] - original_coefficients = ufl.algorithms.extract_coefficients(original_expr) - for coeff in coefficients: - original_coefficient_positions.append(original_coefficients.index(coeff)) - - ir["coefficient_names"] = [ - object_names.get(id(obj), f"w{j}") for j, obj in enumerate(coefficients) - ] - ir["constant_names"] = [ object_names.get(id(obj), f"c{j}") for j, obj in enumerate(ufl.algorithms.analysis.extract_constants(expr)) @@ -641,20 +687,39 @@ def _compute_expression_ir( if len(argument_elements) > 1: raise RuntimeError("Expression with more than one Argument not implemented.") - ir["original_coefficient_positions"] = original_coefficient_positions - - coefficient_elements = tuple(f.ufl_element() for f in coefficients) + coefficients = ufl.algorithms.extract_coefficients(expr) + original_coefficients = ufl.algorithms.extract_coefficients(original_expr) - offsets = {} - _offset = 0 - for i, el in enumerate(coefficient_elements): - offsets[coefficients[i]] = _offset - _offset += element_dimensions[el] + coefficient_numbering: dict[ufl.Coefficient, int] = {} + coefficient_offsets: dict[ufl.Coefficient, int] = {} + proxy_coefficient_numbering: dict[ProxyCoefficient, int] = {} + proxy_coefficient_offsets: dict[ProxyCoefficient, int] = {} + _offset_c = 0 + _offset_p = 0 + original_coefficient_positions = [] + for i, coeff in enumerate(coefficients): + el = coeff.ufl_element() + if isinstance(coeff, ProxyCoefficient): + proxy_coefficient_numbering[coeff] = i + proxy_coefficient_offsets[coeff] = _offset_p + _offset_p += element_dimensions[el] + else: + original_coefficient_positions.append(original_coefficients.index(coeff)) + coefficient_numbering[coeff] = i + coefficient_offsets[coeff] = _offset_c + _offset_c += element_dimensions[el] - # Copy offsets also into IR - base_ir["coefficient_offsets"] = offsets + ir["original_coefficient_positions"] = original_coefficient_positions + base_ir["coefficient_numbering"] = coefficient_numbering + base_ir["coefficient_offsets"] = coefficient_offsets + base_ir["proxy_coefficient_numbering"] = proxy_coefficient_numbering + base_ir["proxy_coefficient_offsets"] = proxy_coefficient_offsets base_ir["integral_type"] = "expression" + ir["coefficient_names"] = [ + object_names.get(id(obj), f"w{j}") for j, obj in enumerate(coefficients) + ] + if cell is not None: if (tdim := cell.topological_dimension) == (pdim := points.shape[1]): base_ir["entity_type"] = "cell" @@ -686,14 +751,15 @@ def _compute_expression_ir( weights = np.array([1.0] * points.shape[0]) rule = QuadratureRule(points, weights) - integrands = {"": {rule: expr}} + integrands: dict[basix.CellType | str, dict[QuadratureRule, ufl.core.expr.Expr]] = { + "": {rule: expr} + } if cell is None: assert ( len(ir["original_coefficient_positions"]) == 0 and len(base_ir["original_constant_offsets"]) == 0 ) - expression_ir = compute_integral_ir( cell, base_ir["integral_type"], diff --git a/ffcx/naming.py b/ffcx/naming.py index e12ca3667..9998d1610 100644 --- a/ffcx/naming.py +++ b/ffcx/naming.py @@ -9,7 +9,7 @@ import hashlib import numbers -from typing import Literal +from typing import Any, Literal import numpy as np import numpy.typing as npt @@ -19,6 +19,36 @@ import ffcx.codegeneration +def extract_expression_renumbering(expr: ufl.core.expr.Expr) -> dict[Any, int]: + """Given a UFL expression, extract the renumbering of all quantities. + + This includes coefficients, constants, arguments and domains. + """ + # FIXME Move this to UFL, cache the computation + coeffs = ufl.algorithms.extract_coefficients(expr) + consts = ufl.algorithms.analysis.extract_constants(expr) + args = ufl.algorithms.analysis.extract_arguments(expr) + + rn = dict() + rn.update(dict((c, i) for i, c in enumerate(coeffs))) + rn.update(dict((c, i) for i, c in enumerate(consts))) + rn.update(dict((c, i) for i, c in enumerate(args))) + + domains: list[ufl.AbstractDomain] = [] + for coeff in coeffs: + domains.append(*ufl.domain.extract_domains(coeff)) + for arg in args: + domains.append(*ufl.domain.extract_domains(arg)) + for gc in ufl.algorithms.analysis.extract_type(expr, ufl.classes.GeometricQuantity): + domains.append(*ufl.domain.extract_domains(gc)) + for const in consts: + domains.append(*ufl.domain.extract_domains(const)) + domains = ufl.algorithms.analysis.unique_tuple(domains) + assert all([isinstance(domain, ufl.Mesh) for domain in domains]) + rn.update(dict((d, i) for i, d in enumerate(domains))) + return rn + + def compute_signature( ufl_objects: list[ufl.Form] | list[tuple[ufl.core.expr.Expr, npt.NDArray[np.floating]]], tag: str, @@ -33,38 +63,26 @@ def compute_signature( if isinstance(ufl_object, ufl.Form): kind = "form" object_signature += ufl_object.signature() + # As UFL doesn't give the correct signature of the form for interpolate + # https://github.com/FEniCS/ufl/issues/473 + # We check for proxycoefficients and append the spaces to the hash + + baseform_ops = ufl.algorithms.extract_base_form_operators(ufl_object) + for op in baseform_ops: + if isinstance(op, ufl.Interpolate): + rn = extract_expression_renumbering(op.ufl_operands[0]) + space_signature = str(op.ufl_function_space()._ufl_signature_data_(rn)) + object_signature += hashlib.sha1(space_signature.encode("utf-8")).hexdigest() + elif isinstance(ufl_object, tuple) and isinstance(ufl_object[0], ufl.core.expr.Expr): expr = ufl_object[0] points = ufl_object[1] - # FIXME Move this to UFL, cache the computation - coeffs = ufl.algorithms.extract_coefficients(expr) - consts = ufl.algorithms.analysis.extract_constants(expr) - args = ufl.algorithms.analysis.extract_arguments(expr) - - rn = dict() - rn.update(dict((c, i) for i, c in enumerate(coeffs))) - rn.update(dict((c, i) for i, c in enumerate(consts))) - rn.update(dict((c, i) for i, c in enumerate(args))) - - domains: list[ufl.AbstractDomain] = [] - for coeff in coeffs: - domains.append(*ufl.domain.extract_domains(coeff)) - for arg in args: - domains.append(*ufl.domain.extract_domains(arg)) - for gc in ufl.algorithms.analysis.extract_type(expr, ufl.classes.GeometricQuantity): - domains.append(*ufl.domain.extract_domains(gc)) - for const in consts: - domains.append(*ufl.domain.extract_domains(const)) - domains = ufl.algorithms.analysis.unique_tuple(domains) - assert all([isinstance(domain, ufl.Mesh) for domain in domains]) - rn.update(dict((d, i) for i, d in enumerate(domains))) - # Hash on UFL signature and points + rn = extract_expression_renumbering(expr) signature = ufl.algorithms.signature.compute_expression_signature(expr, rn) object_signature += signature object_signature += repr(points) - kind = "expression" else: raise RuntimeError(f"Unknown ufl object type {ufl_object.__class__.__name__}") diff --git a/test/test_interpolate.py b/test/test_interpolate.py new file mode 100644 index 000000000..1bd5f34ac --- /dev/null +++ b/test/test_interpolate.py @@ -0,0 +1,99 @@ +import basix.ufl +import numpy as np +import pytest +import ufl + +import ffcx.codegeneration.jit +from ffcx.codegeneration.utils import dtype_to_c_type, dtype_to_scalar_dtype + + +@pytest.mark.parametrize("dtype", ["float64", "float32"]) +@pytest.mark.parametrize("element", [("N1curl", {}), ("Lagrange", {"shape": (2,)})]) +def test_interpolate(compile_args, dtype, element): + + cell = "triangle" + family, el_kwargs = element + + domain = ufl.Mesh(basix.ufl.element("Lagrange", cell, 1, shape=(2,))) + element = basix.ufl.element(family, cell, 2, **el_kwargs) + V_int = ufl.FunctionSpace(domain, element) + + # Space containing other coefficients + Q = ufl.FunctionSpace(domain, basix.ufl.element("Lagrange", cell, 2)) + w = ufl.Coefficient(Q) + z = ufl.Coefficient(Q) + q = ufl.Coefficient(Q) + + x = ufl.SpatialCoordinate(domain) + c = ufl.Constant(domain) + f = ufl.as_vector((z, -x[0])) + ufl.as_vector((q, q)) + If = ufl.Interpolate(f, V_int) + J = (w * If[0] + If[1]) * ufl.dx + f_ref = ufl.as_vector((x[1], -x[0])) + ufl.as_vector((1, 1)) + J_ref = (w * f_ref[0] + f_ref[1]) * ufl.dx + + compiled_forms, module, _code = ffcx.codegeneration.jit.compile_forms( + [J, J_ref], + options={"scalar_type": dtype}, + # cache_dir=".ffcx_cache", + cffi_extra_compile_args=[], # compile_args, + visualise=False, + ) + + xdtype = dtype_to_scalar_dtype(dtype) + scale = 4.2 + coords = np.array([[0.0, 0.0, 0.0], [2.0, 0.0, 0.0], [0.0, scale, 0.0]], dtype=xdtype).flatten() + c = np.array([2.3], dtype=dtype) + # Coefficients are ordered according to when they were created, thus + # w, z, q + # We set z to x[1], w, q to 1 and w to a some non-zero value + q_size = Q.ufl_element().basix_element.dim + assert q_size == 6 + d = np.empty(3 * q_size, dtype=dtype) + d[0:q_size] = [0.2, 0.3, 0.4, 0.5, 0.6, 0.7] # w + d[q_size : 2 * q_size] = [0.0, 0.0, scale * 1.0, 0.5 * scale, 0.5 * scale, 0.0] # z + d[2 * q_size : 3 * q_size] = [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] # q + + # Get kernel + ffi = module.ffi + form = compiled_forms[0] + offsets = form.form_integral_offsets + cell = module.lib.cell + assert offsets[cell + 1] - offsets[cell] == 1 + + default_integral = form.form_integrals[offsets[cell]] + + A = np.zeros(1, dtype=dtype) + c_type, c_xtype = dtype_to_c_type(dtype), dtype_to_c_type(xdtype) + kernel = getattr(default_integral, f"tabulate_tensor_{dtype}") + kernel( + ffi.cast(f"{c_type} *", A.ctypes.data), + ffi.cast(f"{c_type} *", d.ctypes.data), + ffi.cast(f"{c_type} *", c.ctypes.data), + ffi.cast(f"{c_xtype} *", coords.ctypes.data), + ffi.NULL, + ffi.NULL, + ffi.NULL, + ) + + d_ref = np.zeros(q_size, dtype=dtype) + d_ref[:] = d[0:q_size] # w + c_ref = c.copy() + form_ref = compiled_forms[1] + offsets_ref = form_ref.form_integral_offsets + assert offsets_ref[cell + 1] - offsets_ref[cell] == 1 + ref_integral = form_ref.form_integrals[offsets_ref[cell]] + + A_ref = np.zeros(1, dtype=dtype) + ref_kernel = getattr(ref_integral, f"tabulate_tensor_{dtype}") + ref_kernel( + ffi.cast(f"{c_type} *", A_ref.ctypes.data), + ffi.cast(f"{c_type} *", d_ref.ctypes.data), + ffi.cast(f"{c_type} *", c_ref.ctypes.data), + ffi.cast(f"{c_xtype} *", coords.ctypes.data), + ffi.NULL, + ffi.NULL, + ffi.NULL, + ) + tol = np.finfo(dtype).eps * 100 + np.testing.assert_allclose(A, A_ref, atol=tol)