Reduce kernel duplication

.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@dokken/group_integrals
           python3 -m pip install git+https://github.com/FEniCS/basix.git
 
       - name: Install FFCx

.github/workflows/pythonapp.yml

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

ffcx/codegeneration/form.py

@@ -9,6 +9,8 @@
 
 import logging
 
+import numpy
+
 from ffcx.codegeneration import form_template
 
 logger = logging.getLogger("ffcx")
@@ -33,7 +35,10 @@ def generator(ir, parameters):
     code = []
     cases = []
     for itg_type in ("cell", "interior_facet", "exterior_facet"):
-        cases += [(L.Symbol(itg_type), L.Return(len(ir.subdomain_ids[itg_type])))]
+        num_integrals = 0
+        for ids in ir.subdomain_ids[itg_type].values():
+            num_integrals += len(ids)
+        cases += [(L.Symbol(itg_type), L.Return(num_integrals))]
     code += [L.Switch("integral_type", cases, default=L.Return(0))]
     d["num_integrals"] = L.StatementList(code)
 
@@ -68,8 +73,8 @@ def generator(ir, parameters):
     if len(ir.finite_elements) > 0:
         d["finite_elements"] = f"finite_elements_{ir.name}"
         d["finite_elements_init"] = L.ArrayDecl("ufcx_finite_element*", f"finite_elements_{ir.name}", values=[
-                                                L.AddressOf(L.Symbol(el)) for el in ir.finite_elements],
-                                                sizes=len(ir.finite_elements))
+            L.AddressOf(L.Symbol(el)) for el in ir.finite_elements],
+            sizes=len(ir.finite_elements))
     else:
         d["finite_elements"] = L.Null()
         d["finite_elements_init"] = ""
@@ -87,16 +92,26 @@ def generator(ir, parameters):
     code_ids = []
     cases_ids = []
     for itg_type in ("cell", "interior_facet", "exterior_facet"):
-        if len(ir.integral_names[itg_type]) > 0:
+        # Get list of integrals and subdomain_ids for each kernel
+        values = []
+        id_values = []
+        for idx in ir.integral_names[itg_type].keys():
+            integrals = list(ir.subdomain_ids[itg_type][idx])
+            values.extend([L.AddressOf(L.Symbol(ir.integral_names[itg_type][idx]))] * len(integrals))
+            id_values.extend(integrals)
+        value_sort = numpy.argsort(id_values)
+        values = [values[value_sort[i]] for i in range(len(values))]
+        id_values = [id_values[value_sort[i]] for i in range(len(id_values))]
+        if len(values) > 0:
             code += [L.ArrayDecl(
                 "static ufcx_integral*", f"integrals_{itg_type}_{ir.name}",
-                values=[L.AddressOf(L.Symbol(itg)) for itg in ir.integral_names[itg_type]],
-                sizes=len(ir.integral_names[itg_type]))]
+                values=values, sizes=len(values))]
             cases.append((L.Symbol(itg_type), L.Return(L.Symbol(f"integrals_{itg_type}_{ir.name}"))))
 
+        if len(id_values) > 0:
             code_ids += [L.ArrayDecl(
                 "static int", f"integral_ids_{itg_type}_{ir.name}",
-                values=ir.subdomain_ids[itg_type], sizes=len(ir.subdomain_ids[itg_type]))]
+                values=id_values, sizes=len(id_values))]
             cases_ids.append((L.Symbol(itg_type), L.Return(L.Symbol(f"integral_ids_{itg_type}_{ir.name}"))))
 
     code += [L.Switch("integral_type", cases, default=L.Return(L.Null()))]

ffcx/ir/representation.py

@@ -18,12 +18,14 @@
 
 import itertools
 import logging
+import typing
 import warnings
 from collections import namedtuple
 
 import numpy
 import ufl
 from ffcx import naming
+from ffcx.analysis import ufl_data
 from ffcx.element_interface import create_element
 from ffcx.ir.integral import compute_integral_ir
 from ffcx.ir.representationutils import (QuadratureRule,
@@ -63,7 +65,8 @@
 ir_data = namedtuple('ir_data', ['elements', 'dofmaps', 'integrals', 'forms', 'expressions'])
 
 
-def compute_ir(analysis, object_names, prefix, parameters, visualise):
+def compute_ir(analysis: ufl_data, object_names: typing.Dict, prefix: typing.Optional[str], parameters: typing.Dict,
+               visualise: bool):
     """Compute intermediate representation."""
     logger.info(79 * "*")
     logger.info("Compiler stage 2: Computing intermediate representation of objects")
@@ -79,8 +82,9 @@ def compute_ir(analysis, object_names, prefix, parameters, visualise):
     for fd_index, fd in enumerate(analysis.form_data):
         form_names[fd_index] = naming.form_name(fd.original_form, fd_index, prefix)
         for itg_index, itg_data in enumerate(fd.integral_data):
-            integral_names[(fd_index, itg_index)] = naming.integral_name(fd.original_form, itg_data.integral_type,
-                                                                         fd_index, itg_data.subdomain_id, prefix)
+            # Unique ID for integrals is 2**form_index * 3**integral_index
+            integral_names[(fd_index, itg_index)] = naming.integral_name(
+                fd.original_form, itg_data.integral_type, fd_index, itg_index, prefix)
 
     ir_elements = [
         _compute_element_ir(e, analysis.element_numbers, finite_element_names)
@@ -147,11 +151,11 @@ def _compute_element_ir(ufl_element, element_numbers, finite_element_names):
     ir["num_sub_elements"] = ufl_element.num_sub_elements()
     ir["sub_elements"] = [finite_element_names[e] for e in ufl_element.sub_elements()]
 
-    if hasattr(basix_element, "block_size"):
+    try:
         ir["block_size"] = basix_element.block_size
         ufl_element = ufl_element.sub_elements()[0]
         basix_element = create_element(ufl_element)
-    else:
+    except AttributeError:
         ir["block_size"] = 1
 
     ir["entity_dofs"] = basix_element.entity_dofs
@@ -175,10 +179,10 @@ def _compute_dofmap_ir(ufl_element, element_numbers, dofmap_names):
     ir["sub_dofmaps"] = [dofmap_names[e] for e in ufl_element.sub_elements()]
     ir["num_sub_dofmaps"] = ufl_element.num_sub_elements()
 
-    if hasattr(basix_element, "block_size"):
+    try:
         ir["block_size"] = basix_element.block_size
         basix_element = basix_element.sub_element
-    else:
+    except AttributeError:
         ir["block_size"] = 1
 
     # Precompute repeatedly used items
@@ -226,7 +230,6 @@ def _compute_integral_ir(form_data, form_index, element_numbers, integral_names,
         cellname = cell.cellname()
         tdim = cell.topological_dimension()
         assert all(tdim == itg.ufl_domain().topological_dimension() for itg in itg_data.integrals)
-
         ir = {
             "integral_type": itg_data.integral_type,
             "subdomain_id": itg_data.subdomain_id,
@@ -430,32 +433,19 @@ def _compute_form_ir(form_data, form_id, prefix, form_names, integral_names, ele
     # Store names of integrals and subdomain_ids for this form, grouped by integral types
     # Since form points to all integrals it contains, it has to know their names
     # for codegen phase
-    ir["integral_names"] = {}
-    ir["subdomain_ids"] = {}
     ufcx_integral_types = ("cell", "exterior_facet", "interior_facet")
-    for integral_type in ufcx_integral_types:
-        ir["subdomain_ids"][integral_type] = []
-        ir["integral_names"][integral_type] = []
-
-        for itg_index, itg_data in enumerate(form_data.integral_data):
-            if (itg_data.integral_type == integral_type):
-                if itg_data.subdomain_id == "otherwise":
-                    # UFL is using "otherwise" for default integrals (over whole mesh)
-                    # but FFCx needs integers, so otherwise = -1
-                    if len(ir["subdomain_ids"][integral_type]) > 0 and ir["subdomain_ids"][integral_type][0] == -1:
-                        raise ValueError("Only one default ('otherwise') integral allowed.")
-
-                    # Put default integral as first
-                    ir["subdomain_ids"][integral_type] = [-1] + ir["subdomain_ids"][integral_type]
-                    ir["integral_names"][integral_type] = [
-                        integral_names[(form_id, itg_index)]] + ir["integral_names"][integral_type]
-                elif itg_data.subdomain_id < 0:
-                    raise ValueError("Integral subdomain ID must be non-negative.")
-                else:
-                    assert isinstance(itg_data.subdomain_id, int)
-                    ir["subdomain_ids"][integral_type] += [itg_data.subdomain_id]
-                    ir["integral_names"][integral_type] += [integral_names[(form_id, itg_index)]]
+    ir["subdomain_ids"] = {itg_type: {} for itg_type in ufcx_integral_types}
+    ir["integral_names"] = {itg_type: {} for itg_type in ufcx_integral_types}
 
+    for itg_index, itg_data in enumerate(form_data.integral_data):
+        # UFL is using "otherwise" for default integrals (over whole mesh)
+        # but FFCx needs integers, so otherwise = -1
+        integral_type = itg_data.integral_type
+        subdomain_ids = set(sid if sid != "otherwise" else -1 for sid in itg_data.subdomain_id)
+        if min(subdomain_ids) < -1:
+            raise ValueError("Integral subdomain IDs must be non-negative.")
+        ir["subdomain_ids"][integral_type][itg_index] = subdomain_ids
+        ir["integral_names"][integral_type][itg_index] = integral_names[(form_id, itg_index)]
     return ir_form(**ir)
 
 

ffcx/naming.py

@@ -5,17 +5,15 @@
 # SPDX-License-Identifier:    LGPL-3.0-or-later
 
 import hashlib
-
 import ufl
-
+import typing
 import ffcx
 
 
 def compute_signature(ufl_objects, tag):
     """Compute the signature hash.
 
-    Based on the UFL type of the objects and an additional optional
-    'tag'.
+    Based on the UFL type of the objects and an additional optional 'tag'.
     """
     object_signature = ""
     for ufl_object in ufl_objects:
@@ -66,35 +64,37 @@ def compute_signature(ufl_objects, tag):
     return hashlib.sha1(string.encode('utf-8')).hexdigest()
 
 
-def integral_name(original_form, integral_type, form_id, subdomain_id, prefix):
-    sig = compute_signature([original_form], str((prefix, integral_type, form_id, subdomain_id)))
+def integral_name(original_form: ufl.form.Form, integral_type: str, form_id: int, integral_id: int,
+                  prefix: typing.Optional[str]) -> str:
+    """Compute signature for an integral in an ufl form"""
+    sig = compute_signature([original_form], str((prefix, integral_type, form_id, integral_id)))
     return f"integral_{sig}"
 
 
-def form_name(original_form, form_id, prefix):
+def form_name(original_form: ufl.form.Form, form_id: int, prefix: typing.Optional[str]) -> str:
     sig = compute_signature([original_form], str((prefix, form_id)))
     return f"form_{sig}"
 
 
-def finite_element_name(ufl_element, prefix):
+def finite_element_name(ufl_element: ufl.FiniteElementBase, prefix: typing.Optional[str]) -> str:
     assert isinstance(ufl_element, ufl.FiniteElementBase)
     sig = compute_signature([ufl_element], prefix)
     return f"element_{sig}"
 
 
-def dofmap_name(ufl_element, prefix):
+def dofmap_name(ufl_element: ufl.FiniteElementBase, prefix: typing.Optional[str]) -> str:
     assert isinstance(ufl_element, ufl.FiniteElementBase)
     sig = compute_signature([ufl_element], prefix)
     return f"dofmap_{sig}"
 
 
-def expression_name(expression, prefix):
+def expression_name(expression: ufl.core.expr.Expr, prefix: typing.Optional[str]) -> str:
     assert isinstance(expression[0], ufl.core.expr.Expr)
     sig = compute_signature([expression], prefix)
     return f"expression_{sig}"
 
 
-def cdtype_to_numpy(cdtype):
+def cdtype_to_numpy(cdtype: str) -> str:
     """Map a C data type string NumPy datatype string."""
     if cdtype == "double":
         return "float64"

test/test_jit_expression.py

@@ -174,7 +174,6 @@ def test_elimiate_zero_tables_tensor(compile_args):
 
     # Using same basix element for coordinate element and coefficient
     coeff_points = basix_c_e.points
-
     # Compile expression at interpolation points of second order Lagrange space
     b_el = basix.create_element(basix.ElementFamily.P, basix.cell.string_to_type(cell), 0, True)
     points = b_el.points