Remove generated kernel duplication

.github/workflows/dolfin-tests.yml

@@ -46,7 +46,7 @@ jobs:
       - name: Install UFL and Basix (default branches/tags)
         if: github.event_name != 'workflow_dispatch'
         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 UFL and Basix (specified branches/tags)
         if: github.event_name == 'workflow_dispatch'

.github/workflows/pythonapp.yml

@@ -48,7 +48,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/C/form.py

@@ -3,12 +3,16 @@
 # This file is part of FFCx.(https://www.fenicsproject.org)
 #
 # SPDX-License-Identifier:    LGPL-3.0-or-later
-
+#
+# Modified by Chris Richardson and Jørgen S. Dokken 2023
+#
 # Note: Most of the code in this file is a direct translation from the
 # old implementation in FFC
 
 import logging
 
+import numpy
+
 from ffcx.codegeneration.C import form_template
 
 logger = logging.getLogger("ffcx")
@@ -81,8 +85,16 @@ def generator(ir, options):
     integral_offsets = [0]
     # Note: the order of this list is defined by the enum ufcx_integral_type in ufcx.h
     for itg_type in ("cell", "exterior_facet", "interior_facet"):
-        integrals += [f"&{itg}" for itg in ir.integral_names[itg_type]]
-        integral_ids += ir.subdomain_ids[itg_type]
+        unsorted_integrals = []
+        unsorted_ids = []
+        for name, id in zip(ir.integral_names[itg_type], ir.subdomain_ids[itg_type]):
+            unsorted_integrals += [f"&{name}"]
+            unsorted_ids += [id]
+
+        id_sort = numpy.argsort(unsorted_ids)
+        integrals += [unsorted_integrals[i] for i in id_sort]
+        integral_ids += [unsorted_ids[i] for i in id_sort]
+
         integral_offsets.append(len(integrals))
 
     if len(integrals) > 0:

ffcx/ir/representation.py

@@ -18,7 +18,6 @@
 
 import itertools
 import logging
-import numbers
 import typing
 import warnings
 
@@ -562,34 +561,24 @@ def _compute_form_ir(form_data, form_id, prefix, form_names, integral_names, ele
     ir["name_from_uflfile"] = f"form_{prefix}_{form_name}"
 
     # Store names of integrals and subdomain_ids for this form, grouped
-    # by integral types Since form points to all integrals it contains,
+    # 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, numbers.Integral)
-                    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 = [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] += subdomain_ids
+        for _ in range(len(subdomain_ids)):
+            ir["integral_names"][integral_type] += [integral_names[(form_id, itg_index)]]
 
     return FormIR(**ir)
 

test/test_jit_forms.py

@@ -879,3 +879,29 @@ def test_manifold_derivatives(compile_args):
            ffi.cast('uint8_t *', perm.ctypes.data))
 
     assert np.isclose(J[0], 0.0)
+
+
+def test_integral_grouping(compile_args):
+    """
+    We group integrals with common integrands to avoid duplicated integration kernels.
+    This means that `inner(u, v)*dx((1,2,3))  + inner(grad(u), grad(v))*dx(2) + inner(u,v)*dx`
+    is grouped as
+    1. `inner(u,v)*dx(("everywhere", 1, 3))`
+    2. `(inner(grad(u), grad(v)) + inner(u, v))*dx(2)`
+    Each of the forms has one generated `tabulate_tensor_*` function, which is referred to multiple times in
+    `integrals_` and `integral_ids_`
+    """
+    mesh = ufl.Mesh(ufl.VectorElement("Lagrange", ufl.triangle, 1))
+    V = ufl.FunctionSpace(mesh, ufl.FiniteElement("Lagrange", ufl.triangle, 1))
+    u = ufl.TrialFunction(V)
+    v = ufl.TestFunction(V)
+    a = ufl.inner(u, v) * ufl.dx((1, 2, 3)) + ufl.inner(ufl.grad(u), ufl.grad(v)) * ufl.dx(2) + ufl.inner(u, v) * ufl.dx
+    compiled_forms, module, _ = ffcx.codegeneration.jit.compile_forms(
+        [a], cffi_extra_compile_args=compile_args)
+    # NOTE: This assumes that the first integral type is cell integrals, see UFCx.h
+    cell = module.lib.cell
+    num_integrals = compiled_forms[0].form_integral_offsets[cell + 1] - compiled_forms[0].form_integral_offsets[cell]
+    assert num_integrals == 4
+    unique_integrals = set([compiled_forms[0].form_integrals[compiled_forms[0].form_integral_offsets[cell] + i]
+                            for i in range(num_integrals)])
+    assert len(unique_integrals) == 2