Merge branch 'main' into mscroggs/make_quadrature

ffcx/codegeneration/dofmap.py

@@ -15,8 +15,11 @@
 logger = logging.getLogger("ffcx")
 
 
-def tabulate_entity_dofs(L, entity_dofs: typing.List[typing.List[typing.List[int]]],
-                         num_dofs_per_entity: typing.List[int]):
+def tabulate_entity_dofs(
+    L,
+    entity_dofs: typing.List[typing.List[typing.List[int]]],
+    num_dofs_per_entity: typing.List[int],
+):
     # Output argument array
     dofs = L.Symbol("dofs")
 
@@ -30,7 +33,6 @@ def tabulate_entity_dofs(L, entity_dofs: typing.List[typing.List[typing.List[int
     # Generate cases for each dimension:
     all_cases = []
     for dim in range(tdim + 1):
-
         # Ignore if no entities for this dimension
         if num_dofs_per_entity[dim] == 0:
             continue
@@ -39,7 +41,7 @@ def tabulate_entity_dofs(L, entity_dofs: typing.List[typing.List[typing.List[int
         cases = []
         for entity in range(len(entity_dofs[dim])):
             casebody = []
-            for (j, dof) in enumerate(entity_dofs[dim][entity]):
+            for j, dof in enumerate(entity_dofs[dim][entity]):
                 casebody += [L.Assign(dofs[j], dof)]
             cases.append((entity, L.StatementList(casebody)))
 
@@ -64,47 +66,108 @@ def generator(ir, options):
 
     # Attributes
     d["factory_name"] = ir.name
-    d["signature"] = f"\"{ir.signature}\""
+    d["signature"] = f'"{ir.signature}"'
     d["num_global_support_dofs"] = ir.num_global_support_dofs
     d["num_element_support_dofs"] = ir.num_element_support_dofs
     d["num_sub_dofmaps"] = ir.num_sub_dofmaps
 
     import ffcx.codegeneration.C.cnodes as L
 
     num_entity_dofs = ir.num_entity_dofs + [0, 0, 0, 0]
-    num_entity_dofs = num_entity_dofs[: 4]
+    num_entity_dofs = num_entity_dofs[:4]
     d["num_entity_dofs"] = f"num_entity_dofs_{ir.name}"
-    d["num_entity_dofs_init"] = L.ArrayDecl("int", f"num_entity_dofs_{ir.name}",
-                                            values=num_entity_dofs, sizes=4)
+    d["num_entity_dofs_init"] = L.ArrayDecl(
+        "int", f"num_entity_dofs_{ir.name}", values=num_entity_dofs, sizes=4
+    )
 
     num_entity_closure_dofs = ir.num_entity_closure_dofs + [0, 0, 0, 0]
     num_entity_closure_dofs = num_entity_closure_dofs[:4]
     d["num_entity_closure_dofs"] = f"num_entity_closure_dofs_{ir.name}"
-    d["num_entity_closure_dofs_init"] = L.ArrayDecl("int", f"num_entity_closure_dofs_{ir.name}",
-                                                    values=num_entity_closure_dofs, sizes=4)
+    d["num_entity_closure_dofs_init"] = L.ArrayDecl(
+        "int",
+        f"num_entity_closure_dofs_{ir.name}",
+        values=num_entity_closure_dofs,
+        sizes=4,
+    )
+
+    flattened_entity_dofs = []
+    entity_dof_offsets = [0]
+    for dim in ir.entity_dofs:
+        for ent in dim:
+            for v in ent:
+                flattened_entity_dofs.append(v)
+            entity_dof_offsets.append(len(flattened_entity_dofs))
+    d["entity_dofs"] = f"entity_dofs_{ir.name}"
+    d["entity_dofs_init"] = L.ArrayDecl(
+        "int",
+        f"entity_dofs_{ir.name}",
+        values=flattened_entity_dofs,
+        sizes=len(flattened_entity_dofs),
+    )
+    d["entity_dof_offsets"] = f"entity_dof_offsets_{ir.name}"
+    d["entity_dof_offsets_init"] = L.ArrayDecl(
+        "int",
+        f"entity_dof_offsets_{ir.name}",
+        values=entity_dof_offsets,
+        sizes=len(entity_dof_offsets),
+    )
+
+    # Closure
+    flattened_entity_closure_dofs = []
+    entity_closure_dof_offsets = [0]
+    for dim in ir.entity_closure_dofs:
+        for ent in dim:
+            for v in ent:
+                flattened_entity_closure_dofs.append(v)
+            entity_closure_dof_offsets.append(len(flattened_entity_closure_dofs))
+    d["entity_closure_dofs"] = f"entity_closure_dofs_{ir.name}"
+    d["entity_closure_dofs_init"] = L.ArrayDecl(
+        "int",
+        f"entity_closure_dofs_{ir.name}",
+        values=flattened_entity_closure_dofs,
+        sizes=len(flattened_entity_closure_dofs),
+    )
+    d["entity_closure_dof_offsets"] = f"entity_closure_dof_offsets_{ir.name}"
+    d["entity_closure_dof_offsets_init"] = L.ArrayDecl(
+        "int",
+        f"entity_closure_dof_offsets_{ir.name}",
+        values=entity_closure_dof_offsets,
+        sizes=len(entity_closure_dof_offsets),
+    )
 
     d["block_size"] = ir.block_size
 
     # Functions
-    d["tabulate_entity_dofs"] = tabulate_entity_dofs(L, ir.entity_dofs, ir.num_entity_dofs)
-    d["tabulate_entity_closure_dofs"] = tabulate_entity_dofs(L, ir.entity_closure_dofs, ir.num_entity_closure_dofs)
+    d["tabulate_entity_dofs"] = tabulate_entity_dofs(
+        L, ir.entity_dofs, ir.num_entity_dofs
+    )
+    d["tabulate_entity_closure_dofs"] = tabulate_entity_dofs(
+        L, ir.entity_closure_dofs, ir.num_entity_closure_dofs
+    )
 
     if len(ir.sub_dofmaps) > 0:
         d["sub_dofmaps_initialization"] = L.ArrayDecl(
-            "ufcx_dofmap*", f"sub_dofmaps_{ir.name}",
-            values=[L.AddressOf(L.Symbol(dofmap)) for dofmap in ir.sub_dofmaps], sizes=len(ir.sub_dofmaps))
+            "ufcx_dofmap*",
+            f"sub_dofmaps_{ir.name}",
+            values=[L.AddressOf(L.Symbol(dofmap)) for dofmap in ir.sub_dofmaps],
+            sizes=len(ir.sub_dofmaps),
+        )
         d["sub_dofmaps"] = f"sub_dofmaps_{ir.name}"
     else:
         d["sub_dofmaps_initialization"] = ""
         d["sub_dofmaps"] = "NULL"
 
     # Check that no keys are redundant or have been missed
     from string import Formatter
-    fields = [fname for _, fname, _, _ in Formatter().parse(ufcx_dofmap.factory) if fname]
+
+    fields = [
+        fname for _, fname, _, _ in Formatter().parse(ufcx_dofmap.factory) if fname
+    ]
     # Remove square brackets from any field names
     fields = [f.split("[")[0] for f in fields]
     assert set(fields) == set(
-        d.keys()), "Mismatch between keys in template and in formatting dict."
+        d.keys()
+    ), "Mismatch between keys in template and in formatting dict."
 
     # Format implementation code
     implementation = ufcx_dofmap.factory.format_map(d)

ffcx/codegeneration/dofmap_template.py

@@ -26,12 +26,24 @@
 
 {num_entity_closure_dofs_init}
 
+{entity_dofs_init}
+
+{entity_dof_offsets_init}
+
+{entity_closure_dofs_init}
+
+{entity_closure_dof_offsets_init}
+
 ufcx_dofmap {factory_name} =
 {{
   .signature = {signature},
   .num_global_support_dofs = {num_global_support_dofs},
   .num_element_support_dofs = {num_element_support_dofs},
   .block_size = {block_size},
+  .entity_dofs = {entity_dofs},
+  .entity_dof_offsets = {entity_dof_offsets},
+  .entity_closure_dofs = {entity_closure_dofs},
+  .entity_closure_dof_offsets = {entity_closure_dof_offsets},
   .num_entity_dofs = {num_entity_dofs},
   .tabulate_entity_dofs = tabulate_entity_dofs_{factory_name},
   .num_entity_closure_dofs = {num_entity_closure_dofs},

ffcx/codegeneration/form.py

@@ -25,7 +25,7 @@ def generator(ir, options):
     d = {}
     d["factory_name"] = ir.name
     d["name_from_uflfile"] = ir.name_from_uflfile
-    d["signature"] = f"\"{ir.signature}\""
+    d["signature"] = f'"{ir.signature}"'
     d["rank"] = ir.rank
     d["num_coefficients"] = ir.num_coefficients
     d["num_constants"] = ir.num_constants
@@ -39,8 +39,11 @@ def generator(ir, options):
 
     if len(ir.original_coefficient_position) > 0:
         d["original_coefficient_position_init"] = L.ArrayDecl(
-            "int", f"original_coefficient_position_{ir.name}",
-            values=ir.original_coefficient_position, sizes=len(ir.original_coefficient_position))
+            "int",
+            f"original_coefficient_position_{ir.name}",
+            values=ir.original_coefficient_position,
+            sizes=len(ir.original_coefficient_position),
+        )
         d["original_coefficient_position"] = f"original_coefficient_position_{ir.name}"
     else:
         d["original_coefficient_position_init"] = ""
@@ -49,7 +52,7 @@ def generator(ir, options):
     cnames = ir.coefficient_names
     assert ir.num_coefficients == len(cnames)
     names = L.Symbol("names")
-    if (len(cnames) == 0):
+    if len(cnames) == 0:
         code = [L.Return(L.Null())]
     else:
         code = [L.ArrayDecl("static const char*", names, len(cnames), cnames)]
@@ -67,37 +70,102 @@ def generator(ir, options):
 
     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))
+        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),
+        )
     else:
         d["finite_elements"] = L.Null()
         d["finite_elements_init"] = ""
 
     if len(ir.dofmaps) > 0:
         d["dofmaps"] = f"dofmaps_{ir.name}"
-        d["dofmaps_init"] = L.ArrayDecl("ufcx_dofmap*", f"dofmaps_{ir.name}", values=[
-            L.AddressOf(L.Symbol(dofmap)) for dofmap in ir.dofmaps], sizes=len(ir.dofmaps))
+        d["dofmaps_init"] = L.ArrayDecl(
+            "ufcx_dofmap*",
+            f"dofmaps_{ir.name}",
+            values=[L.AddressOf(L.Symbol(dofmap)) for dofmap in ir.dofmaps],
+            sizes=len(ir.dofmaps),
+        )
     else:
         d["dofmaps"] = L.Null()
         d["dofmaps_init"] = ""
 
+    integrals = []
+    integral_ids = []
+    integral_offsets = [0]
+    for itg_type in ("cell", "interior_facet", "exterior_facet"):
+        integrals += [L.AddressOf(L.Symbol(itg)) for itg in ir.integral_names[itg_type]]
+        integral_ids += ir.subdomain_ids[itg_type]
+        integral_offsets.append(len(integrals))
+
+    if len(integrals) > 0:
+        d["form_integrals_init"] = L.ArrayDecl(
+            "static ufcx_integral*",
+            f"form_integrals_{ir.name}",
+            values=integrals,
+            sizes=len(integrals),
+        )
+        d["form_integrals"] = f"form_integrals_{ir.name}"
+        d["form_integral_ids_init"] = L.ArrayDecl(
+            "int",
+            f"form_integral_ids_{ir.name}",
+            values=integral_ids,
+            sizes=len(integral_ids),
+        )
+        d["form_integral_ids"] = f"form_integral_ids_{ir.name}"
+    else:
+        d["form_integrals_init"] = ""
+        d["form_integrals"] = "NULL"
+        d["form_integral_ids_init"] = ""
+        d["form_integral_ids"] = "NULL"
+
+    d["form_integral_offsets_init"] = L.ArrayDecl(
+        "int",
+        f"form_integral_offsets_{ir.name}",
+        values=integral_offsets,
+        sizes=len(integral_offsets),
+    )
+
     code = []
     cases = []
     code_ids = []
     cases_ids = []
     for itg_type in ("cell", "interior_facet", "exterior_facet"):
         if len(ir.integral_names[itg_type]) > 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]))]
-            cases.append((L.Symbol(itg_type), L.Return(L.Symbol(f"integrals_{itg_type}_{ir.name}"))))
-
-            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]))]
-            cases_ids.append((L.Symbol(itg_type), L.Return(L.Symbol(f"integral_ids_{itg_type}_{ir.name}"))))
+            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]),
+                )
+            ]
+            cases.append(
+                (
+                    L.Symbol(itg_type),
+                    L.Return(L.Symbol(f"integrals_{itg_type}_{ir.name}")),
+                )
+            )
+
+            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]),
+                )
+            ]
+            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()))]
     code_ids += [L.Switch("integral_type", cases_ids, default=L.Return(L.Null()))]
@@ -110,12 +178,19 @@ def generator(ir, options):
 
     # FIXME: Should be handled differently, revise how
     # ufcx_function_space is generated
-    for (name, (element, dofmap, cmap_family, cmap_degree, cmap_celltype, cmap_variant)) in ir.function_spaces.items():
+    for name, (
+        element,
+        dofmap,
+        cmap_family,
+        cmap_degree,
+        cmap_celltype,
+        cmap_variant,
+    ) in ir.function_spaces.items():
         code += [f"static ufcx_function_space functionspace_{name} ="]
         code += ["{"]
         code += [f".finite_element = &{element},"]
         code += [f".dofmap = &{dofmap},"]
-        code += [f".geometry_family = \"{cmap_family}\","]
+        code += [f'.geometry_family = "{cmap_family}",']
         code += [f".geometry_degree = {cmap_degree},"]
         code += [f".geometry_basix_cell = {int(cmap_celltype)},"]
         code += [f".geometry_basix_variant = {int(cmap_variant)}"]
@@ -133,14 +208,20 @@ def generator(ir, options):
 
     # Check that no keys are redundant or have been missed
     from string import Formatter
-    fields = [fname for _, fname, _, _ in Formatter().parse(form_template.factory) if fname]
-    assert set(fields) == set(d.keys()), "Mismatch between keys in template and in formatting dict"
+
+    fields = [
+        fname for _, fname, _, _ in Formatter().parse(form_template.factory) if fname
+    ]
+    assert set(fields) == set(
+        d.keys()
+    ), "Mismatch between keys in template and in formatting dict"
 
     # Format implementation code
     implementation = form_template.factory.format_map(d)
 
     # Format declaration
-    declaration = form_template.declaration.format(factory_name=d["factory_name"],
-                                                   name_from_uflfile=d["name_from_uflfile"])
+    declaration = form_template.declaration.format(
+        factory_name=d["factory_name"], name_from_uflfile=d["name_from_uflfile"]
+    )
 
     return declaration, implementation

ffcx/codegeneration/form_template.py

@@ -24,6 +24,9 @@
 {original_coefficient_position_init}
 {dofmaps_init}
 {finite_elements_init}
+{form_integral_offsets_init}
+{form_integrals_init}
+{form_integral_ids_init}
 
 // Return a list of the coefficient names.
 const char** coefficient_name_{factory_name}(void)
@@ -70,7 +73,11 @@
   .integral_ids = integral_ids_{factory_name},
   .num_integrals = num_integrals_{factory_name},
 
-  .integrals = integrals_{factory_name}
+  .integrals = integrals_{factory_name},
+
+  .form_integrals = {form_integrals},
+  .form_integral_ids = {form_integral_ids},
+  .form_integral_offsets = form_integral_offsets_{factory_name}
 }};
 
 // Alias name