Update and extend ufc_expression

.github/workflows/dolfin-tests.yml

@@ -33,7 +33,7 @@ jobs:
           python3 -m pip install --upgrade pip
       - name: Install UFL and Basix
         run: |
-          python3 -m pip install git+https://github.com/FEniCS/ufl.git
+          python3 -m pip install git+https://github.com/FEniCS/ufl.git@michal/expressions
           python3 -m pip install git+https://github.com/FEniCS/basix.git
 
       - name: Install FFCx
@@ -45,7 +45,7 @@ jobs:
         with:
           path: ./dolfinx
           repository: FEniCS/dolfinx
-          ref: main
+          ref: michal/ufc-expression
       - name: Clone and install DOLFINx
         run: |
           cmake -G Ninja -DCMAKE_BUILD_TYPE=Developer -B build -S dolfinx/cpp/

.github/workflows/pythonapp.yml

@@ -42,7 +42,8 @@ jobs:
 
       - name: Install FEniCS dependencies (Python)
         run: |
-          python -m pip install git+https://github.com/FEniCS/ufl.git
+          python -m pip install --upgrade pip setuptools
+          python -m pip install git+https://github.com/FEniCS/ufl.git@michal/expressions
           python -m pip install git+https://github.com/FEniCS/basix.git
 
       - name: Install FFCx

demo/GradExpression.ufl

@@ -0,0 +1,36 @@
+# Copyright (C) 2021 Michal Habera
+#
+# This file is part of FFCX.
+#
+# FFC is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# FFC is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with FFC. If not, see <http://www.gnu.org/licenses/>.
+#
+# Defines an Expression which evaluates gradient of an Argument
+# at predefined set of points.
+#
+# Compile this form with FFC: ffcx GradExpression.ufl
+
+element = FiniteElement("Lagrange", triangle, 2)
+coordinate_el = VectorElement("Lagrange", triangle, 1)
+mesh = Mesh(coordinate_el)
+
+V = FunctionSpace(mesh, element)
+
+u = TrialFunction(V)
+f = Coefficient(V)
+c = Constant(mesh)
+
+grad_u = c * f * grad(u)
+
+expressions = [(grad_u, [[0.0, 0.1], [0.2, 0.3]]),
+               (c * f * u, [[0.0, 1.0], [1.0, 0.0]])]

ffcx/analysis.py

@@ -25,8 +25,7 @@
                                    'unique_coordinate_elements', 'expressions'])
 
 
-def analyze_ufl_objects(ufl_objects: typing.Union[typing.List[ufl.form.Form], typing.List[ufl.FiniteElement],
-                                                  typing.List],
+def analyze_ufl_objects(ufl_objects: typing.List,
                         parameters: typing.Dict) -> ufl_data:
     """Analyze ufl object(s).
 
@@ -51,62 +50,51 @@ def analyze_ufl_objects(ufl_objects: typing.Union[typing.List[ufl.form.Form], ty
     logger.info("Compiler stage 1: Analyzing UFL objects")
     logger.info(79 * "*")
 
-    unique_elements = set()
-    unique_coordinate_elements = set()
+    elements = []
+    coordinate_elements = []
+
+    # Group objects by types
+    forms = []
     expressions = []
+    processed_expressions = []
+
+    for ufl_object in ufl_objects:
+        if isinstance(ufl_object, ufl.form.Form):
+            forms += [ufl_object]
+        elif isinstance(ufl_object, ufl.FiniteElementBase):
+            elements += [ufl_object]
+        elif isinstance(ufl_object, ufl.Mesh):
+            coordinate_elements += [ufl_object.ufl_coordinate_element()]
+        elif isinstance(ufl_object[0], ufl.core.expr.Expr):
+            original_expression = ufl_object[0]
+            points = numpy.asarray(ufl_object[1])
+            expressions += [(original_expression, points)]
+        else:
+            raise TypeError("UFL objects not recognised.")
 
-    # FIXME: This assumes that forms come before elements in
-    # ufl_objects? Is this reasonable?
-    if isinstance(ufl_objects[0], ufl.form.Form):
-        forms = ufl_objects
-        form_data = tuple(_analyze_form(form, parameters) for form in forms)
-
-        # Extract unique elements across forms
-        for data in form_data:
-            unique_elements.update(data.unique_sub_elements)
-
-        # Extract uniquecoordinate elements across forms
-        for data in form_data:
-            unique_coordinate_elements.update(data.coordinate_elements)
-    elif isinstance(ufl_objects[0], ufl.FiniteElementBase):
-        form_data = ()
-        # Extract unique (sub)elements
-        elements = ufl_objects
-        unique_elements.update(ufl.algorithms.analysis.extract_sub_elements(elements))
-    elif isinstance(ufl_objects[0], ufl.Mesh):
-        form_data = ()
-        # Extract unique (sub)elements
-        meshes = ufl_objects
-        unique_coordinate_elements = set(mesh.ufl_coordinate_element() for mesh in meshes)
-    elif isinstance(ufl_objects[0], tuple) and isinstance(ufl_objects[0][0], ufl.core.expr.Expr):
-        form_data = ()
-        for expression in ufl_objects:
-            original_expression = expression[0]
-            points = expression[1]
-            expression = expression[0]
-
-            unique_elements.update(ufl.algorithms.extract_elements(expression))
-            unique_elements.update(ufl.algorithms.extract_sub_elements(unique_elements))
-
-            expression = _analyze_expression(expression, parameters)
-            expressions.append((expression, points, original_expression))
-    else:
-        raise TypeError("UFL objects not recognised.")
+    form_data = tuple(_analyze_form(form, parameters) for form in forms)
+    for data in form_data:
+        elements += data.unique_sub_elements
+        coordinate_elements += data.coordinate_elements
+
+    for original_expression, points in expressions:
+        elements += list(ufl.algorithms.extract_elements(original_expression))
+        processed_expression = _analyze_expression(original_expression, parameters)
+        processed_expressions += [(processed_expression, points, original_expression)]
 
-    # Make sure coordinate elements and their subelements are included
-    unique_elements.update(ufl.algorithms.analysis.extract_sub_elements(unique_coordinate_elements))
+    elements += ufl.algorithms.analysis.extract_sub_elements(elements)
 
     # Sort elements so sub-elements come before mixed elements
-    unique_elements = ufl.algorithms.sort_elements(unique_elements)
-    unique_coordinate_element_list = sorted(unique_coordinate_elements, key=lambda x: repr(x))
+    unique_elements = ufl.algorithms.sort_elements(set(elements))
+    unique_coordinate_element_list = sorted(set(coordinate_elements), key=lambda x: repr(x))
 
     # Compute dict (map) from element to index
     element_numbers = {element: i for i, element in enumerate(unique_elements)}
 
     return ufl_data(form_data=form_data, unique_elements=unique_elements,
                     element_numbers=element_numbers,
                     unique_coordinate_elements=unique_coordinate_element_list,
-                    expressions=expressions)
+                    expressions=processed_expressions)
 
 
 def _analyze_expression(expression: ufl.core.expr.Expr, parameters: typing.Dict):

ffcx/codegeneration/C/format_lines.py

@@ -39,7 +39,7 @@ def iter_indented_lines(snippets, level=0):
     - tuple,list: Yield lines from recursive application of this function to list items.
 
     """
-    tabsize = 4
+    tabsize = 2
     indentation = ' ' * (tabsize * level)
     if isinstance(snippets, str):
         for line in snippets.split("\n"):

ffcx/codegeneration/expressions.py

@@ -10,11 +10,14 @@
 from itertools import product
 
 import ufl
+
+from ffcx.codegeneration import geometry
 from ffcx.codegeneration import expressions_template
 from ffcx.codegeneration.backend import FFCXBackend
 from ffcx.codegeneration.C.format_lines import format_indented_lines
 from ffcx.codegeneration.C.cnodes import CNode
 from ffcx.ir.representation import ir_expression
+from ffcx.naming import cdtype_to_numpy
 
 logger = logging.getLogger("ffcx")
 
@@ -28,14 +31,17 @@ def generator(ir, parameters):
     factory_name = ir.name
 
     # Format declaration
-    declaration = expressions_template.declaration.format(factory_name=factory_name)
+    declaration = expressions_template.declaration.format(
+        factory_name=factory_name, name_from_uflfile=ir.name_from_uflfile)
 
     backend = FFCXBackend(ir, parameters)
     L = backend.language
     eg = ExpressionGenerator(ir, backend)
 
     d = {}
+    d["name_from_uflfile"] = ir.name_from_uflfile
     d["factory_name"] = ir.name
+
     parts = eg.generate()
 
     body = format_indented_lines(parts.cs_format(), 1)
@@ -64,10 +70,56 @@ def generator(ir, parameters):
 
     d["num_components"] = len(ir.expression_shape)
     d["num_coefficients"] = len(ir.coefficient_numbering)
+    d["num_constants"] = len(ir.constant_names)
     d["num_points"] = ir.points.shape[0]
     d["topological_dimension"] = ir.points.shape[1]
-    d["needs_facet_permutations"] = "true" if ir.needs_facet_permutations else "false"
     d["scalar_type"] = parameters["scalar_type"]
+    d["np_scalar_type"] = cdtype_to_numpy(parameters["scalar_type"])
+    d["rank"] = len(ir.tensor_shape)
+
+    if len(ir.coefficient_names) > 0:
+        d["coefficient_names_init"] = L.ArrayDecl(
+            "static const char*", f"coefficient_names_{ir.name}", values=ir.coefficient_names,
+            sizes=len(ir.coefficient_names))
+        d["coefficient_names"] = f"coefficient_names_{ir.name}"
+    else:
+        d["coefficient_names_init"] = ""
+        d["coefficient_names"] = L.Null()
+
+    if len(ir.constant_names) > 0:
+        d["constant_names_init"] = L.ArrayDecl(
+            "static const char*", f"constant_names_{ir.name}", values=ir.coefficient_names,
+            sizes=len(ir.coefficient_names))
+        d["constant_names"] = f"constant_names_{ir.name}"
+    else:
+        d["constant_names_init"] = ""
+        d["constant_names"] = L.Null()
+
+    code = []
+
+    # FIXME: Should be handled differently, revise how
+    # ufcx_function_space is generated (also for ufcx_form)
+    for (name, (element, dofmap, cmap_family, cmap_degree)) in ir.function_spaces.items():
+        code += [f"static ufcx_function_space function_space_{name}_{ir.name_from_uflfile} ="]
+        code += ["{"]
+        code += [f".finite_element = &{element},"]
+        code += [f".dofmap = &{dofmap},"]
+        code += [f".geometry_family = \"{cmap_family}\","]
+        code += [f".geometry_degree = {cmap_degree}"]
+        code += ["};"]
+
+    d["function_spaces_alloc"] = L.StatementList(code)
+    d["function_spaces"] = ""
+
+    if len(ir.function_spaces) > 0:
+        d["function_spaces"] = f"function_spaces_{ir.name}"
+        d["function_spaces_init"] = L.ArrayDecl("ufcx_function_space*", f"function_spaces_{ir.name}", values=[
+                                                L.AddressOf(L.Symbol(f"function_space_{name}_{ir.name_from_uflfile}"))
+                                                for (name, _) in ir.function_spaces.items()],
+                                                sizes=len(ir.function_spaces))
+    else:
+        d["function_spaces"] = L.Null()
+        d["function_spaces_init"] = ""
 
     # Check that no keys are redundant or have been missed
     from string import Formatter
@@ -101,6 +153,8 @@ def generate(self):
         parts = []
 
         parts += self.generate_element_tables()
+        # Generate the tables of geometry data that are needed
+        parts += self.generate_geometry_tables()
         parts += self.generate_piecewise_partition()
 
         all_preparts = []
@@ -116,6 +170,31 @@ def generate(self):
 
         return L.StatementList(parts)
 
+    def generate_geometry_tables(self):
+        """Generate static tables of geometry data."""
+        L = self.backend.language
+
+        # Currently we only support circumradius
+        ufl_geometry = {
+            ufl.geometry.ReferenceCellVolume: "reference_cell_volume"
+        }
+        cells = {t: set() for t in ufl_geometry.keys()}
+
+        for integrand in self.ir.integrand.values():
+            for attr in integrand["factorization"].nodes.values():
+                mt = attr.get("mt")
+                if mt is not None:
+                    t = type(mt.terminal)
+                    if t in ufl_geometry:
+                        cells[t].add(mt.terminal.ufl_domain().ufl_cell().cellname())
+
+        parts = []
+        for i, cell_list in cells.items():
+            for c in cell_list:
+                parts.append(geometry.write_table(L, ufl_geometry[i], c))
+
+        return parts
+
     def generate_element_tables(self):
         """Generate tables of FE basis evaluated at specified points."""
         L = self.backend.language

ffcx/codegeneration/expressions_template.py

@@ -6,35 +6,55 @@
 
 declaration = """
 extern ufcx_expression {factory_name};
+
+// Helper used to create expression using name which was given to the
+// expression in the UFL file.
+// This helper is called in user c++ code.
+//
+extern ufcx_expression* {name_from_uflfile};
 """
 
 factory = """
 // Code for expression {factory_name}
 
-void tabulate_expression_{factory_name}({scalar_type}* restrict A,
-                                        const {scalar_type}* restrict w,
-                                        const {scalar_type}* restrict c,
-                                        const double* restrict coordinate_dofs)
+void tabulate_tensor_{factory_name}({scalar_type}* restrict A,
+                                    const {scalar_type}* restrict w,
+                                    const {scalar_type}* restrict c,
+                                    const double* restrict coordinate_dofs,
+                                    const int* restrict entity_local_index,
+                                    const uint8_t* restrict quadrature_permutation)
 {{
 {tabulate_expression}
 }}
 
 {points_init}
 {value_shape_init}
 {original_coefficient_positions_init}
+{function_spaces_alloc}
+{function_spaces_init}
+{coefficient_names_init}
+{constant_names_init}
+
 
 ufcx_expression {factory_name} =
 {{
-  .tabulate_expression = tabulate_expression_{factory_name},
+  .tabulate_tensor_{np_scalar_type} = 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},
   .topological_dimension = {topological_dimension},
-  .needs_facet_permutations = {needs_facet_permutations},
   .points = {points},
   .value_shape = {value_shape},
   .num_components = {num_components},
-  .original_coefficient_positions = {original_coefficient_positions}
+  .rank = {rank},
+  .function_spaces = {function_spaces}
 }};
 
+// Alias name
+ufcx_expression* {name_from_uflfile} = &{factory_name};
+
 // End of code for expression {factory_name}
 """