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

.github/workflows/pythonapp.yml

@@ -43,7 +43,7 @@ jobs:
       - name: Install FEniCS dependencies (Python)
         run: |
           python -m pip install --upgrade pip setuptools
-          python -m pip install git+https://github.com/FEniCS/ufl.git
+          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,34 @@
+# 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)
+
+expressions = [(c * f * grad(u), [[0.0, 0.1], [0.2, 0.3]]),
+               (c * f * u, [[0.0, 1.0], [1.0, 0.0]])]

ffcx/analysis.py

@@ -79,10 +79,12 @@ def analyze_ufl_objects(ufl_objects: typing.Union[typing.List[ufl.form.Form], ty
         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 = ()
+        if not all(isinstance(expression[0], ufl.core.expr.Expr) for expression in ufl_objects):
+            raise RuntimeError("All objects in the list must be UFL Expressions")
+
         for expression in ufl_objects:
             original_expression = expression[0]
-            points = expression[1]
+            points = numpy.asarray(expression[1])
             expression = expression[0]
 
             unique_elements.update(ufl.algorithms.extract_elements(expression))

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

@@ -15,6 +15,7 @@
 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")
 
@@ -68,6 +69,16 @@ def generator(ir, parameters):
     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.tensor_shape) > 0:
+        d["num_argument_dofs_init"] = L.ArrayDecl(
+            "static int", f"num_argument_dofs_{ir.name}", values=ir.tensor_shape, sizes=len(ir.tensor_shape))
+        d["num_argument_dofs"] = f"num_argument_dofs_{ir.name}"
+    else:
+        d["num_argument_dofs_init"] = ""
+        d["num_argument_dofs"] = L.Null()
 
     # Check that no keys are redundant or have been missed
     from string import Formatter

ffcx/codegeneration/expressions_template.py

@@ -11,29 +11,34 @@
 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}
+{num_argument_dofs_init}
 
 ufc_expression {factory_name} =
 {{
-  .tabulate_expression = tabulate_expression_{factory_name},
+  .tabulate_tensor_{np_scalar_type} = tabulate_tensor_{factory_name},
   .num_coefficients = {num_coefficients},
   .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}
+  .original_coefficient_positions = {original_coefficient_positions},
+  .rank = {rank},
+  .num_argument_dofs = {num_argument_dofs}
 }};
 
 // End of code for expression {factory_name}

ffcx/codegeneration/ufc.h

@@ -1,7 +1,7 @@
 /// This is UFC (Unified Form-assembly Code)
 /// This code is released into the public domain.
 ///
-/// The FEniCS Project (http://www.fenicsproject.org/) 2006-2019.
+/// The FEniCS Project (http://www.fenicsproject.org/) 2006-2021.
 ///
 /// UFC defines the interface between code generated by FFCx and the
 /// DOLFINx C++ library. Changes here must be reflected both in the FFCx
@@ -272,29 +272,37 @@ extern "C"
 
   typedef struct ufc_expression
   {
-
-    /// Evaluate expression into tensor A with compiled evaluation
-    /// points
+    /// Evaluate expression into tensor A with compiled evaluation points
     ///
     /// @param[out] A
-    /// @param[in] w Coefficients attached to the expression.
-    /// Dimensions: w[coefficient][dof].
-    /// @param[in] c Constants attached to the expression. Dimensions:
-    /// c[constant][dim].
-    /// @param[in] coordinate_dofs Values of degrees of freedom of
-    /// coordinate element. Defines the geometry of the cell.
-    /// Dimensions: coordinate_dofs[num_dofs][3].
-    void (*tabulate_expression)(double* restrict A,
-                                const double* restrict w,
-                                const double* restrict c,
-                                const double* restrict coordinate_dofs);
-
-    /// Positions of coefficients in original expression
-    const int* original_coefficient_positions;
+    ///         Dimensions: A[num_points][num_components][num_argument_dofs]
+    ///
+    /// @see ufc_tabulate_tensor
+    ///
+    ufc_tabulate_tensor_float32* tabulate_tensor_float32;
+    ufc_tabulate_tensor_float64* tabulate_tensor_float64;
+    ufc_tabulate_tensor_longdouble* tabulate_tensor_longdouble;
+    ufc_tabulate_tensor_complex64* tabulate_tensor_complex64;
+    ufc_tabulate_tensor_complex128* tabulate_tensor_complex128;
+
+    /// String identifying the expression
+    const char* signature;
 
     /// Number of coefficients
     int num_coefficients;
 
+    /// Number of constants
+    int num_constants;
+
+    /// Original coefficient position for each coefficient
+    const int* original_coefficient_positions;
+
+    /// List of names of coefficients
+    const char** coefficient_names;
+
+    /// List of names of constants
+    const char** constant_names;
+
     /// Number of evaluation points
     int num_points;
 
@@ -314,6 +322,14 @@ extern "C"
 
     /// Number of components of return_shape
     int num_components;
+
+    /// Rank, i.e. number of arguments
+    int rank;
+
+    /// Number of dofs for Arguments
+    /// Dimensions: num_argument_dofs[rank].
+    const int* num_argument_dofs;
+
   } ufc_expression;
 
   /// This class defines the interface for the assembly of the global

ffcx/ir/representation.py

@@ -57,7 +57,7 @@
     'name', 'element_dimensions', 'params', 'unique_tables', 'unique_table_types', 'integrand',
     'table_dofmaps', 'coefficient_numbering', 'coefficient_offsets',
     'integral_type', 'entitytype', 'tensor_shape', 'expression_shape', 'original_constant_offsets',
-    'original_coefficient_positions', 'points', 'needs_facet_permutations'])
+    'original_coefficient_positions', 'points', 'needs_facet_permutations', 'coefficient_names', 'constant_names'])
 
 ir_data = namedtuple('ir_data', ['elements', 'dofmaps', 'integrals', 'forms', 'expressions'])
 
@@ -104,7 +104,7 @@ def compute_ir(analysis, object_names, prefix, parameters, visualise):
         for (i, fd) in enumerate(analysis.form_data)
     ]
 
-    ir_expressions = [_compute_expression_ir(expr, i, prefix, analysis, parameters, visualise)
+    ir_expressions = [_compute_expression_ir(expr, i, prefix, analysis, parameters, visualise, object_names)
                       for i, expr in enumerate(analysis.expressions)]
 
     return ir_data(elements=ir_elements, dofmaps=ir_dofmaps,
@@ -388,10 +388,10 @@ def _compute_form_ir(form_data, form_id, prefix, form_names, integral_names, ele
     ir["num_coefficients"] = len(form_data.reduced_coefficients)
     ir["num_constants"] = len(form_data.original_form.constants())
 
-    ir["coefficient_names"] = [object_names.get(id(obj), "w%d" % j)
+    ir["coefficient_names"] = [object_names.get(id(obj), f"w{j}")
                                for j, obj in enumerate(form_data.reduced_coefficients)]
 
-    ir["constant_names"] = [object_names.get(id(obj), "c%d" % j)
+    ir["constant_names"] = [object_names.get(id(obj), f"c{j}")
                             for j, obj in enumerate(form_data.original_form.constants())]
 
     ir["original_coefficient_position"] = form_data.original_coefficient_positions
@@ -451,7 +451,7 @@ def _compute_form_ir(form_data, form_id, prefix, form_names, integral_names, ele
     return ir_form(**ir)
 
 
-def _compute_expression_ir(expression, index, prefix, analysis, parameters, visualise):
+def _compute_expression_ir(expression, index, prefix, analysis, parameters, visualise, object_names):
     """Compute intermediate representation of expression."""
     logger.info(f"Computing IR for expression {index}")
 
@@ -505,6 +505,12 @@ def _compute_expression_ir(expression, index, prefix, analysis, parameters, visu
     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(expression))]
+
     ir["original_coefficient_positions"] = original_coefficient_positions
 
     coefficient_elements = tuple(f.ufl_element() for f in coefficients)

ffcx/main.py

@@ -71,6 +71,9 @@ def main(args=None):
         if len(ufd.forms) > 0:
             code_h, code_c = compiler.compile_ufl_objects(
                 ufd.forms, ufd.object_names, prefix=prefix, parameters=parameters, visualise=xargs.visualise)
+        elif len(ufd.expressions) > 0:
+            code_h, code_c = compiler.compile_ufl_objects(
+                ufd.expressions, ufd.object_names, prefix=prefix, parameters=parameters, visualise=xargs.visualise)
         else:
             code_h, code_c = compiler.compile_ufl_objects(
                 ufd.elements, ufd.object_names, prefix=prefix, parameters=parameters, visualise=xargs.visualise)

test/test_jit_expression.py

@@ -60,16 +60,20 @@ def test_matvec(compile_args):
     # Coefficient storage XYXYXY
     w = np.array(f_mat.T.flatten(), dtype=np_type)
     c = np.array([0.5], dtype=np_type)
+    entity_index = np.array([0], dtype=np.intc)
+    quad_perm = np.array([0], dtype=np.dtype("uint8"))
 
     # Coords storage XYZXYZXYZ
     coords = np.array([[0.0, 0.0, 0.0],
                        [1.0, 0.0, 0.0],
                        [0.0, 1.0, 0.0]], dtype=np.float64)
-    expression.tabulate_expression(
+    expression.tabulate_tensor_float64(
         ffi.cast('{type} *'.format(type=c_type), A.ctypes.data),
         ffi.cast('{type} *'.format(type=c_type), w.ctypes.data),
         ffi.cast('{type} *'.format(type=c_type), c.ctypes.data),
-        ffi.cast('double *', coords.ctypes.data))
+        ffi.cast('double *', coords.ctypes.data),
+        ffi.cast('int *', entity_index.ctypes.data),
+        ffi.cast('uint8_t *', quad_perm.ctypes.data))
 
     # Check the computation against correct NumPy value
     assert np.allclose(A, 0.5 * np.dot(a_mat, f_mat).T)
@@ -118,15 +122,19 @@ def test_rank1(compile_args):
     # Coefficient storage XYXYXY
     w = np.array([0.0], dtype=np_type)
     c = np.array([0.0], dtype=np_type)
+    entity_index = np.array([0], dtype=np.intc)
+    quad_perm = np.array([0], dtype=np.dtype("uint8"))
 
     # Coords storage XYZXYZXYZ
     coords = np.zeros((points.shape[0], 3), dtype=np.float64)
     coords[:, :2] = points
-    expression.tabulate_expression(
+    expression.tabulate_tensor_float64(
         ffi.cast('{type} *'.format(type=c_type), A.ctypes.data),
         ffi.cast('{type} *'.format(type=c_type), w.ctypes.data),
         ffi.cast('{type} *'.format(type=c_type), c.ctypes.data),
-        ffi.cast('double *', coords.ctypes.data))
+        ffi.cast('double *', coords.ctypes.data),
+        ffi.cast('int *', entity_index.ctypes.data),
+        ffi.cast('uint8_t *', quad_perm.ctypes.data))
 
     f = np.array([[1.0, 2.0, 3.0], [-4.0, -5.0, 6.0]])