diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 02b8a2ae5b..e6ae7e5aa9 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -1247,8 +1247,12 @@ def iter_active_coefficients(form, kinfo): @staticmethod def iter_constants(form, kinfo): # FIXME kinfo not currently needed from tsfc.ufl_utils import extract_firedrake_constants - for const in extract_firedrake_constants(form): - yield const + if isinstance(form, slate.TensorBase): + for const in form.constants(): + yield const + else: + for const in extract_firedrake_constants(form): + yield const @staticmethod def index_function_spaces(form, indices): diff --git a/firedrake/slate/slac/kernel_builder.py b/firedrake/slate/slac/kernel_builder.py index 3e4e95c183..e94252d2f1 100644 --- a/firedrake/slate/slac/kernel_builder.py +++ b/firedrake/slate/slac/kernel_builder.py @@ -9,6 +9,7 @@ from loopy.symbolic import SubArrayRef import pymbolic.primitives as pym +from firedrake.constant import Constant import firedrake.slate.slate as slate from firedrake.slate.slac.tsfc_driver import compile_terminal_form @@ -107,13 +108,16 @@ def shape(self, tensor): else: return tensor.shape - def extent(self, coefficient): - """ Calculation of the range of a coefficient.""" - element = coefficient.ufl_element() - if element.family() == "Real": - return (coefficient.dat.cdim, ) + def extent(self, argument): + """ Calculation of the range of a constant or coefficient.""" + if isinstance(argument, Constant): + return (argument.dat.cdim, ) else: - return (create_element(element).space_dimension(), ) + element = argument.ufl_element() + if element.family() == "Real": + return (argument.dat.cdim, ) + else: + return (create_element(element).space_dimension(), ) def generate_lhs(self, tensor, temp): """ Generation of an lhs for the loopy kernel, @@ -123,7 +127,7 @@ def generate_lhs(self, tensor, temp): lhs = pym.Subscript(temp, idx) return SubArrayRef(idx, lhs) - def collect_tsfc_kernel_data(self, mesh, tsfc_coefficients, wrapper_coefficients, kinfo): + def collect_tsfc_kernel_data(self, mesh, tsfc_coefficients, tsfc_constants, wrapper_coefficients, wrapper_constants, kinfo): """ Collect the kernel data aka the parameters fed into the subkernel, that are coordinates, orientations, cell sizes and cofficients. """ @@ -147,11 +151,14 @@ def collect_tsfc_kernel_data(self, mesh, tsfc_coefficients, wrapper_coefficients ind, = tsfc_coefficients[c] if ind != 0: raise ValueError(f"Active indices of non-mixed function must be (0, ), not {tsfc_coefficients[c]}") - kernel_data.extend([(c, cinfo[0])]) + kernel_data.append((c, cinfo[0])) else: for ind, (c_, info) in enumerate(cinfo.items()): if ind in tsfc_coefficients[c]: - kernel_data.extend([(c_, info[0])]) + kernel_data.append((c_, info[0])) + + # Pick the constants associated with a Tensor()/TSFC kernel + kernel_data.extend([(c, c.name) for c in wrapper_constants if c in tsfc_constants]) return kernel_data def loopify_tsfc_kernel_data(self, kernel_data): @@ -244,6 +251,10 @@ def collect_coefficients(self): coeff_dict[coeff] = (f"w_{i}", self.extent(coeff)) return coeff_dict + def collect_constants(self): + """ All constants of self.expression as a list """ + return self.expression.constants() + def initialise_terminals(self, var2terminal, coefficients): """ Initilisation of the variables in which coefficients and the Tensors coming from TSFC are saved. @@ -349,6 +360,14 @@ def generate_wrapper_kernel_args(self, tensor2temp): dtype=self.tsfc_parameters["scalar_type"]) args.append(kernel_args.CoefficientKernelArg(coeff_loopy_arg)) + for constant in self.bag.constants: + constant_loopy_arg = loopy.GlobalArg( + constant.name, + shape=constant.dat.cdim, + dtype=self.tsfc_parameters["scalar_type"] + ) + args.append(kernel_args.ConstantKernelArg(constant_loopy_arg)) + if self.bag.needs_cell_facets: # Arg for is exterior (==0)/interior (==1) facet or not facet_loopy_arg = loopy.GlobalArg(self.cell_facets_arg_name, @@ -398,7 +417,14 @@ def generate_tsfc_calls(self, terminal, loopy_tensor): output_var = pym.Variable(loopy_tensor.name) reads.append(output_var) output = self.generate_lhs(slate_tensor, output_var) - kernel_data = self.collect_tsfc_kernel_data(mesh, cxt_kernel.coefficients, self.bag.coefficients, kinfo) + kernel_data = self.collect_tsfc_kernel_data( + mesh, + cxt_kernel.coefficients, + cxt_kernel.constants, + self.bag.coefficients, + self.bag.constants, + kinfo + ) reads.extend(self.loopify_tsfc_kernel_data(kernel_data)) # Generate predicates for different integral types @@ -427,8 +453,9 @@ def generate_tsfc_calls(self, terminal, loopy_tensor): class SlateWrapperBag: - def __init__(self, coeffs): + def __init__(self, coeffs, constants): self.coefficients = coeffs + self.constants = constants self.inames = OrderedDict() self.needs_cell_orientations = False self.needs_cell_sizes = False diff --git a/firedrake/slate/slac/tsfc_driver.py b/firedrake/slate/slac/tsfc_driver.py index c7698d7f4d..34baf10dbc 100644 --- a/firedrake/slate/slac/tsfc_driver.py +++ b/firedrake/slate/slac/tsfc_driver.py @@ -5,6 +5,8 @@ from firedrake.slate.slac.utils import RemoveRestrictions from firedrake.tsfc_interface import compile_form as tsfc_compile +from tsfc.ufl_utils import extract_firedrake_constants + from ufl.algorithms.map_integrands import map_integrand_dags from ufl import Form @@ -12,6 +14,7 @@ ContextKernel = collections.namedtuple("ContextKernel", ["tensor", "coefficients", + "constants", "original_integral_type", "tsfc_kernels"]) ContextKernel.__doc__ = """\ @@ -22,6 +25,7 @@ list of TSFC assembly kernels. :param coefficients: The local coefficients of the tensor contained in the integrands (arguments for TSFC subkernels). +:param constants: The local constants of the tensor contained in the integrands. :param original_integral_type: The unmodified measure type of the form integrals. :param tsfc_kernels: A list of local tensor assembly kernels @@ -65,6 +69,7 @@ def compile_terminal_form(tensor, prefix, *, tsfc_parameters=None): if kernels: cxt_k = ContextKernel(tensor=tensor, coefficients=form.coefficients(), + constants=extract_firedrake_constants(form), original_integral_type=orig_it_type, tsfc_kernels=kernels) cxt_kernels.append(cxt_k) diff --git a/firedrake/slate/slac/utils.py b/firedrake/slate/slac/utils.py index 9195806707..ea14231b9b 100644 --- a/firedrake/slate/slac/utils.py +++ b/firedrake/slate/slac/utils.py @@ -195,7 +195,8 @@ def merge_loopy(slate_loopy, output_arg, builder, var2terminal, name): """ Merges tsfc loopy kernels and slate loopy kernel into a wrapper kernel.""" from firedrake.slate.slac.kernel_builder import SlateWrapperBag coeffs = builder.collect_coefficients() - builder.bag = SlateWrapperBag(coeffs) + constants = builder.collect_constants() + builder.bag = SlateWrapperBag(coeffs, constants) # In the initialisation the loopy tensors for the terminals are generated # Those are the needed again for generating the TSFC calls @@ -224,6 +225,13 @@ def merge_loopy(slate_loopy, output_arg, builder, var2terminal, name): # Add profiling for the whole kernel insns, slate_wrapper_event, preamble = profile_insns(name, insns, PETSc.Log.isActive()) + # Add a no-op touching all kernel arguments to make sure they are not + # silently dropped + noop = lp.CInstruction( + (), "", read_variables=frozenset({a.name for a in loopy_args}), + within_inames=frozenset(), within_inames_is_final=True) + insns.append(noop) + # Inames come from initialisations + loopyfying kernel args and lhs domains = builder.bag.index_creator.domains diff --git a/firedrake/slate/slate.py b/firedrake/slate/slate.py index 7c02014bde..c393b90546 100644 --- a/firedrake/slate/slate.py +++ b/firedrake/slate/slate.py @@ -36,6 +36,8 @@ from firedrake.formmanipulation import ExtractSubBlock +from tsfc.ufl_utils import extract_firedrake_constants + __all__ = ['AssembledVector', 'Block', 'Factorization', 'Tensor', 'Inverse', 'Transpose', 'Negative', @@ -213,6 +215,10 @@ def rank(self): def coefficients(self): """Returns a tuple of coefficients associated with the tensor.""" + @abstractmethod + def constants(self): + """Returns a tuple of constants associated with the tensor.""" + @abstractmethod def slate_coefficients(self): """Returns a tuple of Slate coefficients associated with the tensor.""" @@ -463,6 +469,9 @@ def coefficients(self): """Returns a tuple of coefficients associated with the tensor.""" return (self._function,) + def constants(self): + return () + def slate_coefficients(self): """Returns a tuple of coefficients associated with the tensor.""" return self.coefficients() @@ -707,6 +716,11 @@ def coefficients(self): tensor, = self.operands return tensor.coefficients() + def constants(self): + """Returns a tuple of constants associated with the tensor.""" + tensor, = self.operands + return tensor.constants() + def slate_coefficients(self): """Returns a tuple of coefficients associated with the tensor.""" return self.coefficients() @@ -797,6 +811,11 @@ def coefficients(self): tensor, = self.operands return tensor.coefficients() + def constants(self): + """Returns a tuple of constants associated with the tensor.""" + tensor, = self.operands + return tensor.constants() + def slate_coefficients(self): """Returns a tuple of coefficients associated with the tensor.""" return self.coefficients() @@ -896,6 +915,10 @@ def coefficients(self): """Returns a tuple of coefficients associated with the tensor.""" return self.form.coefficients() + def constants(self): + """Returns a tuple of constants associated with the tensor.""" + return unique(extract_firedrake_constants(self.form)) + def slate_coefficients(self): """Returns a tuple of coefficients associated with the tensor.""" return self.coefficients() @@ -944,6 +967,11 @@ def coefficients(self): coeffs = [op.coefficients() for op in self.operands] return tuple(OrderedDict.fromkeys(chain(*coeffs))) + def constants(self): + """Returns a tuple of constants associated with the tensor.""" + const = [op.constants() for op in self.operands] + return unique(chain(*const)) + def slate_coefficients(self): """Returns the expected coefficients of the resulting tensor.""" coeffs = [op.slate_coefficients() for op in self.operands] @@ -1360,6 +1388,16 @@ def space_equivalence(A, B): return A.mesh() == B.mesh() and A.ufl_element() == B.ufl_element() +def unique(iterable): + """ Return tuple of unique items in iterable, items must be hashable + """ + # Use dict to preserve order and compare by hash + unique_dict = {} + for item in iterable: + unique_dict[item] = None + return tuple(unique_dict.keys()) + + # Establishes levels of precedence for Slate tensors precedences = [ [AssembledVector, Block, Factorization, Tensor, DiagonalTensor, Reciprocal],