From 30ead42d32d4c7dbcd97f62393d2ecac9f63ec7b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= Date: Fri, 11 Aug 2023 08:47:36 +0000 Subject: [PATCH] Try repeating fix from prior to https://github.com/FEniCS/ffcx/pull/587 --- ffcx/codegeneration/form.py | 60 +++++++++++++++---------------------- ffcx/ir/representation.py | 38 ++++++++--------------- 2 files changed, 37 insertions(+), 61 deletions(-) diff --git a/ffcx/codegeneration/form.py b/ffcx/codegeneration/form.py index 6de4c46b9..77c57328f 100644 --- a/ffcx/codegeneration/form.py +++ b/ffcx/codegeneration/form.py @@ -8,7 +8,7 @@ # old implementation in FFC import logging - +import numpy from ffcx.codegeneration import form_template logger = logging.getLogger("ffcx") @@ -96,8 +96,10 @@ def generator(ir, options): 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] + 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] integral_offsets.append(len(integrals)) if len(integrals) > 0: @@ -133,39 +135,25 @@ def generator(ir, options): 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}")), - ) - ) + # 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()))] diff --git a/ffcx/ir/representation.py b/ffcx/ir/representation.py index 5513bdc3e..24175ed78 100644 --- a/ffcx/ir/representation.py +++ b/ffcx/ir/representation.py @@ -562,35 +562,23 @@ 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 = set(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][itg_index] = subdomain_ids + ir["integral_names"][integral_type][itg_index] = integral_names[(form_id, itg_index)] + return FormIR(**ir)