Make ds integrals on prism/pyramids generate kernels for each facet type (FEniCS#739)

* working on quadrature

* working on prism facet integrals

* progress towards ds on prisms

* store tables by integrand

* correct test

* update expressions

* ruff

* mypy

* mypy

* ufcx_vertex

* ruff

* ruff

* more ruff

* don't overwrite cell

* remove ufcx_cell_type enum

* ruff

* use basix cell type throughout rather than converting to/from string

* mypy

* set()

* make demo xfail

* handle "vertex"

* typing

* remove testing ipython embed

* fix merge

Author: mscroggs

ffcx/codegeneration/C/expressions.py

@@ -25,7 +25,7 @@ def generator(ir: ExpressionIR, options):
     """Generate UFC code for an expression."""
     logger.info("Generating code for expression:")
     assert len(ir.expression.integrand) == 1, "Expressions only support single quadrature rule"
-    points = next(iter(ir.expression.integrand)).points
+    points = next(iter(ir.expression.integrand))[1].points
     logger.info(f"--- points: {points}")
     factory_name = ir.expression.name
     logger.info(f"--- name: {factory_name}")

ffcx/codegeneration/C/form.py

@@ -86,29 +86,44 @@ def generator(ir: FormIR, options):
     integrals = []
     integral_ids = []
     integral_offsets = [0]
+    integral_domains = []
     # 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"):
         unsorted_integrals = []
         unsorted_ids = []
-        for name, id in zip(ir.integral_names[itg_type], ir.subdomain_ids[itg_type]):
+        unsorted_domains = []
+        for name, domains, id in zip(
+            ir.integral_names[itg_type],
+            ir.integral_domains[itg_type],
+            ir.subdomain_ids[itg_type],
+        ):
             unsorted_integrals += [f"&{name}"]
             unsorted_ids += [id]
+            unsorted_domains += [domains]
 
         id_sort = np.argsort(unsorted_ids)
         integrals += [unsorted_integrals[i] for i in id_sort]
         integral_ids += [unsorted_ids[i] for i in id_sort]
+        integral_domains += [unsorted_domains[i] for i in id_sort]
 
-        integral_offsets.append(len(integrals))
+        integral_offsets.append(sum(len(d) for d in integral_domains))
 
     if len(integrals) > 0:
-        sizes = len(integrals)
-        values = ", ".join(integrals)
+        sizes = sum(len(domains) for domains in integral_domains)
+        values = ", ".join(
+            [
+                f"{i}_{domain.name}"
+                for i, domains in zip(integrals, integral_domains)
+                for domain in domains
+            ]
+        )
         d["form_integrals_init"] = (
             f"static ufcx_integral* form_integrals_{ir.name}[{sizes}] = {{{values}}};"
         )
         d["form_integrals"] = f"form_integrals_{ir.name}"
-        sizes = len(integral_ids)
-        values = ", ".join(str(i) for i in integral_ids)
+        values = ", ".join(
+            f"{i}" for i, domains in zip(integral_ids, integral_domains) for _ in domains
+        )
         d["form_integral_ids_init"] = f"int form_integral_ids_{ir.name}[{sizes}] = {{{values}}};"
         d["form_integral_ids"] = f"form_integral_ids_{ir.name}"
     else:

ffcx/codegeneration/C/integrals.py

@@ -8,6 +8,7 @@
 import logging
 import sys
 
+import basix
 import numpy as np
 
 from ffcx.codegeneration.backend import FFCXBackend
@@ -20,14 +21,13 @@
 logger = logging.getLogger("ffcx")
 
 
-def generator(ir: IntegralIR, options):
+def generator(ir: IntegralIR, domain: basix.CellType, options):
     """Generate C code for an integral."""
     logger.info("Generating code for integral:")
     logger.info(f"--- type: {ir.expression.integral_type}")
     logger.info(f"--- name: {ir.expression.name}")
 
-    """Generate code for an integral."""
-    factory_name = ir.expression.name
+    factory_name = f"{ir.expression.name}_{domain.name}"
 
     # Format declaration
     declaration = ufcx_integrals.declaration.format(factory_name=factory_name)
@@ -39,7 +39,7 @@ def generator(ir: IntegralIR, options):
     ig = IntegralGenerator(ir, backend)
 
     # Generate code ast for the tabulate_tensor body
-    parts = ig.generate()
+    parts = ig.generate(domain)
 
     # Format code as string
     CF = CFormatter(options["scalar_type"])
@@ -52,9 +52,9 @@ def generator(ir: IntegralIR, options):
         values = ", ".join("1" if i else "0" for i in ir.enabled_coefficients)
         sizes = len(ir.enabled_coefficients)
         code["enabled_coefficients_init"] = (
-            f"bool enabled_coefficients_{ir.expression.name}[{sizes}] = {{{values}}};"
+            f"bool enabled_coefficients_{ir.expression.name}_{domain.name}[{sizes}] = {{{values}}};"
         )
-        code["enabled_coefficients"] = f"enabled_coefficients_{ir.expression.name}"
+        code["enabled_coefficients"] = f"enabled_coefficients_{ir.expression.name}_{domain.name}"
     else:
         code["enabled_coefficients_init"] = ""
         code["enabled_coefficients"] = "NULL"
@@ -88,6 +88,7 @@ def generator(ir: IntegralIR, options):
         tabulate_tensor_float64=code["tabulate_tensor_float64"],
         tabulate_tensor_complex64=code["tabulate_tensor_complex64"],
         tabulate_tensor_complex128=code["tabulate_tensor_complex128"],
+        domain=int(domain),
     )
 
     return declaration, implementation

ffcx/codegeneration/C/integrals_template.py

@@ -32,6 +32,7 @@
   {tabulate_tensor_complex128}
   .needs_facet_permutations = {needs_facet_permutations},
   .coordinate_element_hash = {coordinate_element_hash},
+  .domain = {domain},
 }};
 
 // End of code for integral {factory_name}

ffcx/codegeneration/access.py

@@ -237,7 +237,14 @@ def reference_normal(self, mt, tabledata, access):
     def cell_facet_jacobian(self, mt, tabledata, num_points):
         """Access a cell facet jacobian."""
         cellname = ufl.domain.extract_unique_domain(mt.terminal).ufl_cell().cellname()
-        if cellname in ("triangle", "tetrahedron", "quadrilateral", "hexahedron"):
+        if cellname in (
+            "triangle",
+            "tetrahedron",
+            "quadrilateral",
+            "hexahedron",
+            "prism",
+            "pyramid",
+        ):
             table = L.Symbol(f"{cellname}_cell_facet_jacobian", dtype=L.DataType.REAL)
             facet = self.symbols.entity("facet", mt.restriction)
             return table[facet][mt.component[0]][mt.component[1]]

ffcx/codegeneration/codegeneration.py

@@ -45,7 +45,11 @@ def generate_code(ir: DataIR, options: dict[str, int | float | npt.DTypeLike]) -
     logger.info("Compiler stage 3: Generating code")
     logger.info(79 * "*")
 
-    code_integrals = [integral_generator(integral_ir, options) for integral_ir in ir.integrals]
+    code_integrals = [
+        integral_generator(integral_ir, domain, options)
+        for integral_ir in ir.integrals
+        for domain in set(i[0] for i in integral_ir.expression.integrand.keys())
+    ]
     code_forms = [form_generator(form_ir, options) for form_ir in ir.forms]
     code_expressions = [
         expression_generator(expression_ir, options) for expression_ir in ir.expressions

ffcx/codegeneration/expression_generator.py

@@ -97,7 +97,7 @@ def generate_element_tables(self):
         """Generate tables of FE basis evaluated at specified points."""
         parts = []
 
-        tables = self.ir.expression.unique_tables
+        tables = self.ir.expression.unique_tables[self.quadrature_rule[0]]
         table_names = sorted(tables)
 
         for name in table_names:
@@ -125,7 +125,7 @@ def generate_quadrature_loop(self):
         # Generate varying partition
         body = self.generate_varying_partition()
         body = L.commented_code_list(
-            body, f"Points loop body setup quadrature loop {self.quadrature_rule.id()}"
+            body, f"Points loop body setup quadrature loop {self.quadrature_rule[1].id()}"
         )
 
         # Generate dofblock parts, some of this
@@ -139,7 +139,7 @@ def generate_quadrature_loop(self):
             quadparts = []
         else:
             iq = self.backend.symbols.quadrature_loop_index
-            num_points = self.quadrature_rule.points.shape[0]
+            num_points = self.quadrature_rule[1].points.shape[0]
             quadparts = [L.ForRange(iq, 0, num_points, body=body)]
         return preparts, quadparts
 
@@ -148,11 +148,11 @@ def generate_varying_partition(self):
         # Get annotated graph of factorisation
         F = self.ir.expression.integrand[self.quadrature_rule]["factorization"]
 
-        arraysymbol = L.Symbol(f"sv_{self.quadrature_rule.id()}", dtype=L.DataType.SCALAR)
+        arraysymbol = L.Symbol(f"sv_{self.quadrature_rule[1].id()}", dtype=L.DataType.SCALAR)
         parts = self.generate_partition(arraysymbol, F, "varying")
         parts = L.commented_code_list(
             parts,
-            f"Unstructured varying computations for quadrature rule {self.quadrature_rule.id()}",
+            f"Unstructured varying computations for quadrature rule {self.quadrature_rule[1].id()}",
         )
         return parts
 
@@ -216,7 +216,7 @@ def generate_block_parts(self, blockmap, blockdata):
         assert not blockdata.transposed, "Not handled yet"
         components = ufl.product(self.ir.expression.shape)
 
-        num_points = self.quadrature_rule.points.shape[0]
+        num_points = self.quadrature_rule[1].points.shape[0]
         A_shape = [num_points, components] + self.ir.expression.tensor_shape
         A = self.backend.symbols.element_tensor
         iq = self.backend.symbols.quadrature_loop_index