diff --git a/ffcx/codegeneration/C/form.py b/ffcx/codegeneration/C/form.py index 275fe0d8a..30c72091a 100644 --- a/ffcx/codegeneration/C/form.py +++ b/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: diff --git a/ffcx/ir/representation.py b/ffcx/ir/representation.py index fd8158047..fc425391e 100644 --- a/ffcx/ir/representation.py +++ b/ffcx/ir/representation.py @@ -18,7 +18,6 @@ import itertools import logging -import numbers import typing import warnings @@ -560,34 +559,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) diff --git a/test/test_jit_forms.py b/test/test_jit_forms.py index 7ae89ec7c..5ace4f7a7 100644 --- a/test/test_jit_forms.py +++ b/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