Merge in main and add back sorting of integrals based on subdomain id

ffcx/codegeneration/dofmap.py

@@ -8,54 +8,12 @@
 # old implementation in FFC
 
 import logging
-import typing
 
 import ffcx.codegeneration.dofmap_template as ufcx_dofmap
 
 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],
-):
-    # Output argument array
-    dofs = L.Symbol("dofs")
-
-    # Input arguments
-    d = L.Symbol("d")
-    i = L.Symbol("i")
-
-    # TODO: Removed check for (d <= tdim + 1)
-    tdim = len(num_dofs_per_entity) - 1
-
-    # 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
-
-        # Generate cases for each mesh entity
-        cases = []
-        for entity in range(len(entity_dofs[dim])):
-            casebody = []
-            for j, dof in enumerate(entity_dofs[dim][entity]):
-                casebody += [L.Assign(dofs[j], dof)]
-            cases.append((entity, L.StatementList(casebody)))
-
-        # Generate inner switch
-        # TODO: Removed check for (i <= num_entities-1)
-        inner_switch = L.Switch(i, cases, autoscope=False)
-        all_cases.append((dim, inner_switch))
-
-    if all_cases:
-        return L.Switch(d, all_cases, autoscope=False)
-    else:
-        return L.NoOp()
-
-
 def generator(ir, options):
     """Generate UFC code for a dofmap."""
     logger.info("Generating code for dofmap:")
@@ -73,23 +31,6 @@ def generator(ir, options):
 
     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]
-    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
-    )
-
-    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,
-    )
-
     flattened_entity_dofs = []
     entity_dof_offsets = [0]
     for dim in ir.entity_dofs:
@@ -137,14 +78,6 @@ def generator(ir, options):
 
     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
-    )
-
     if len(ir.sub_dofmaps) > 0:
         d["sub_dofmaps_initialization"] = L.ArrayDecl(
             "ufcx_dofmap*",

ffcx/codegeneration/dofmap_template.py

@@ -12,20 +12,6 @@
 
 {sub_dofmaps_initialization}
 
-void tabulate_entity_dofs_{factory_name}(int* restrict dofs, int d, int i)
-{{
-{tabulate_entity_dofs}
-}}
-
-void tabulate_entity_closure_dofs_{factory_name}(int* restrict dofs, int d, int i)
-{{
-{tabulate_entity_closure_dofs}
-}}
-
-{num_entity_dofs_init}
-
-{num_entity_closure_dofs_init}
-
 {entity_dofs_init}
 
 {entity_dof_offsets_init}
@@ -44,10 +30,6 @@
   .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},
-  .tabulate_entity_closure_dofs = tabulate_entity_closure_dofs_{factory_name},
   .num_sub_dofmaps = {num_sub_dofmaps},
   .sub_dofmaps = {sub_dofmaps}
 }};

ffcx/codegeneration/form.py

@@ -30,16 +30,6 @@ def generator(ir, options):
     d["num_coefficients"] = ir.num_coefficients
     d["num_constants"] = ir.num_constants
 
-    code = []
-    cases = []
-    for itg_type in ("cell", "interior_facet", "exterior_facet"):
-        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)
-
     if len(ir.original_coefficient_position) > 0:
         d["original_coefficient_position_init"] = L.ArrayDecl(
             "int",
@@ -98,12 +88,19 @@ def generator(ir, options):
     integrals = []
     integral_ids = []
     integral_offsets = [0]
-    # NOTE: This ordering is dependent of the enums in dolfinx::fem::IntegralType
+    # 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 key, name in ir.integral_names[itg_type].items():
             for subdomain_id in ir.subdomain_ids[itg_type][key]:
-                integrals += [L.AddressOf(L.Symbol(name))]
-                integral_ids += [subdomain_id]
+                unsorted_integrals += [L.AddressOf(L.Symbol(name))]
+                unsorted_ids += [subdomain_id]
+
+        id_sort = numpy.argsort(unsorted_ids)
+        integrals += [unsorted_integrals[id_sort[i]] for i in range(len(unsorted_integrals))]
+        integral_ids += [unsorted_ids[id_sort[i]] for i in range(len(unsorted_integrals))]
+
         integral_offsets.append(len(integrals))
 
     if len(integrals) > 0:
@@ -134,37 +131,6 @@ def generator(ir, options):
         sizes=len(integral_offsets),
     )
 
-    code = []
-    cases = []
-    code_ids = []
-    cases_ids = []
-    for itg_type in ("cell", "interior_facet", "exterior_facet"):
-        # 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=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=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()))]
-    code_ids += [L.Switch("integral_type", cases_ids, default=L.Return(L.Null()))]
-    d["integrals"] = L.StatementList(code)
-
-    d["integral_ids"] = L.StatementList(code_ids)
-
     code = []
     function_name = L.Symbol("function_name")
 

ffcx/codegeneration/form_template.py

@@ -40,21 +40,6 @@
 {constant_name_map}
 }}
 
-int* integral_ids_{factory_name}(ufcx_integral_type integral_type)
-{{
-{integral_ids}
-}}
-
-int num_integrals_{factory_name}(ufcx_integral_type integral_type)
-{{
-{num_integrals}
-}}
-
-ufcx_integral** integrals_{factory_name}(ufcx_integral_type integral_type)
-{{
-{integrals}
-}}
-
 ufcx_form {factory_name} =
 {{
 
@@ -70,11 +55,6 @@
   .finite_elements = {finite_elements},
   .dofmaps = {dofmaps},
 
-  .integral_ids = integral_ids_{factory_name},
-  .num_integrals = num_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}

ffcx/codegeneration/ufcx.h

@@ -456,15 +456,6 @@ extern "C"
     ///        Coefficient number j=i-r if r+j <= i < r+n
     ufcx_dofmap** dofmaps;
 
-    /// All ids for integrals
-    int* (*integral_ids)(ufcx_integral_type);
-
-    /// Number of integrals
-    int (*num_integrals)(ufcx_integral_type);
-
-    /// Get an integral on sub domain subdomain_id
-    ufcx_integral** (*integrals)(ufcx_integral_type);
-
     /// List of cell, interior facet and exterior facet integrals
     ufcx_integral** form_integrals;
 

test/test_add_mode.py

@@ -35,11 +35,13 @@ def test_additive_facet_integral(mode, compile_args):
     ffi = module.ffi
     form0 = compiled_forms[0]
 
-    assert form0.num_integrals(module.lib.exterior_facet) == 1
-    ids = form0.integral_ids(module.lib.exterior_facet)
-    assert ids[0] == -1
+    integral_offsets = form0.form_integral_offsets
+    ex = module.lib.exterior_facet
+    assert integral_offsets[ex + 1] - integral_offsets[ex] == 1
+    integral_id = form0.form_integral_ids[integral_offsets[ex]]
+    assert integral_id == -1
 
-    default_integral = form0.integrals(module.lib.exterior_facet)[0]
+    default_integral = form0.form_integrals[integral_offsets[ex]]
 
     np_type = cdtype_to_numpy(mode)
     A = np.zeros((3, 3), dtype=np_type)
@@ -83,11 +85,14 @@ def test_additive_cell_integral(mode, compile_args):
     ffi = module.ffi
     form0 = compiled_forms[0]
 
-    assert form0.num_integrals(module.lib.cell) == 1
-    ids = form0.integral_ids(module.lib.cell)
-    assert ids[0] == -1
+    cell = module.lib.cell
+    offsets = form0.form_integral_offsets
+    num_integrals = offsets[cell + 1] - offsets[cell]
+    assert num_integrals == 1
+    integral_id = form0.form_integral_ids[offsets[cell]]
+    assert integral_id == -1
 
-    default_integral = form0.integrals(0)[0]
+    default_integral = form0.form_integrals[offsets[cell]]
 
     np_type = cdtype_to_numpy(mode)
     A = np.zeros((3, 3), dtype=np_type)

test/test_blocked_elements.py

@@ -32,15 +32,11 @@ def test_finite_element(compile_args):
     assert ufcx_dofmap.num_global_support_dofs == 0
     assert ufcx_dofmap.num_global_support_dofs == 0
     assert ufcx_dofmap.num_element_support_dofs == 3
-    assert ufcx_dofmap.num_entity_dofs[0] == 1
-    assert ufcx_dofmap.num_entity_dofs[1] == 0
-    assert ufcx_dofmap.num_entity_dofs[2] == 0
-    assert ufcx_dofmap.num_entity_dofs[3] == 0
+    off = np.array([ufcx_dofmap.entity_dof_offsets[i] for i in range(8)])
+    assert np.all(np.diff(off) == [1, 1, 1, 0, 0, 0, 0])
+
     for v in range(3):
-        vals = np.zeros(1, dtype=np.int32)
-        vals_ptr = module.ffi.cast("int *", module.ffi.from_buffer(vals))
-        ufcx_dofmap.tabulate_entity_dofs(vals_ptr, 0, v)
-        assert vals[0] == v
+        assert ufcx_dofmap.entity_dofs[v] == v
     assert ufcx_dofmap.num_sub_dofmaps == 0
 
 
@@ -66,15 +62,11 @@ def test_vector_element(compile_args):
     assert ufcx_dofmap.num_global_support_dofs == 0
     assert ufcx_dofmap.num_global_support_dofs == 0
     assert ufcx_dofmap.num_element_support_dofs == 3
-    assert ufcx_dofmap.num_entity_dofs[0] == 1
-    assert ufcx_dofmap.num_entity_dofs[1] == 0
-    assert ufcx_dofmap.num_entity_dofs[2] == 0
-    assert ufcx_dofmap.num_entity_dofs[3] == 0
+    off = np.array([ufcx_dofmap.entity_dof_offsets[i] for i in range(8)])
+    assert np.all(np.diff(off) == [1, 1, 1, 0, 0, 0, 0])
+
     for v in range(3):
-        vals = np.zeros(1, dtype=np.int32)
-        vals_ptr = module.ffi.cast("int *", module.ffi.from_buffer(vals))
-        ufcx_dofmap.tabulate_entity_dofs(vals_ptr, 0, v)
-        assert vals[0] == v
+        assert ufcx_dofmap.entity_dofs[v] == v
     assert ufcx_dofmap.num_sub_dofmaps == 2
 
 
@@ -102,15 +94,11 @@ def test_tensor_element(compile_args):
     assert ufcx_dofmap.num_global_support_dofs == 0
     assert ufcx_dofmap.num_global_support_dofs == 0
     assert ufcx_dofmap.num_element_support_dofs == 3
-    assert ufcx_dofmap.num_entity_dofs[0] == 1
-    assert ufcx_dofmap.num_entity_dofs[1] == 0
-    assert ufcx_dofmap.num_entity_dofs[2] == 0
-    assert ufcx_dofmap.num_entity_dofs[3] == 0
+    off = np.array([ufcx_dofmap.entity_dof_offsets[i] for i in range(8)])
+    assert np.all(np.diff(off) == [1, 1, 1, 0, 0, 0, 0])
+
     for v in range(3):
-        vals = np.zeros(1, dtype=np.int32)
-        vals_ptr = module.ffi.cast("int *", module.ffi.from_buffer(vals))
-        ufcx_dofmap.tabulate_entity_dofs(vals_ptr, 0, v)
-        assert vals[0] == v
+        assert ufcx_dofmap.entity_dofs[v] == v
     assert ufcx_dofmap.num_sub_dofmaps == 4
 
 
@@ -136,14 +124,10 @@ def test_vector_quadrature_element(compile_args):
     assert ufcx_dofmap.num_global_support_dofs == 0
     assert ufcx_dofmap.num_global_support_dofs == 0
     assert ufcx_dofmap.num_element_support_dofs == 4
-    assert ufcx_dofmap.num_entity_dofs[0] == 0
-    assert ufcx_dofmap.num_entity_dofs[1] == 0
-    assert ufcx_dofmap.num_entity_dofs[2] == 0
-    assert ufcx_dofmap.num_entity_dofs[3] == 4
-
-    vals = np.zeros(4, dtype=np.int32)
-    vals_ptr = module.ffi.cast("int *", module.ffi.from_buffer(vals))
-    ufcx_dofmap.tabulate_entity_dofs(vals_ptr, 3, 0)
-    assert (vals == [0, 1, 2, 3]).all()
+    off = np.array([ufcx_dofmap.entity_dof_offsets[i] for i in range(16)])
+    assert np.all(np.diff(off) == [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 4])
+
+    for i in range(4):
+        assert ufcx_dofmap.entity_dofs[i] == i
 
     assert ufcx_dofmap.num_sub_dofmaps == 3