diff --git a/tsfc/driver.py b/tsfc/driver.py index d6db7120..7ef5b8e1 100644 --- a/tsfc/driver.py +++ b/tsfc/driver.py @@ -1,9 +1,7 @@ import collections -import operator -import string import time import sys -from functools import reduce +from functools import partial from itertools import chain from numpy import asarray, allclose @@ -11,22 +9,19 @@ import ufl from ufl.algorithms import extract_arguments, extract_coefficients from ufl.algorithms.analysis import has_type -from ufl.classes import Form, GeometricQuantity +from ufl.classes import Form, GeometricQuantity, Coefficient, FunctionSpace from ufl.log import GREEN -from ufl.utils.sequences import max_degree import gem import gem.impero_utils as impero_utils import FIAT -from FIAT.reference_element import TensorProductCell from FIAT.functional import PointEvaluation from finat.point_set import PointSet -from finat.quadrature import AbstractQuadratureRule, make_quadrature, QuadratureRule +from finat.quadrature import QuadratureRule from tsfc import fem, ufl_utils -from tsfc.finatinterface import as_fiat_cell from tsfc.logging import logger from tsfc.parameters import default_parameters, is_complex from tsfc.ufl_utils import apply_mapping @@ -35,6 +30,172 @@ sys.setrecursionlimit(3000) +class TSFCFormData(object): + r"""Mimic `ufl.FormData`. + + :arg form_data_tuple: A tuple of `ufl.FormData`s. + :arg extraarg_tuple: A tuple of extra `ufl.Argument`s + corresponding to form_data_tuple. These extra + arguments are eventually replaced by the user with + the associated functions in function_tuple after + compiling UFL but before compiling gem. These + arguments thus do not contribute to the rank of the form. + :arg function_tuple: A tuple of functions corresponding + to extraarg_tuple. + :arg original_form: The form from which forms for + `ufl.Formdata`s were extracted. + :diagonal: A flag for diagonal matrix assembly. + + This class mimics `ufl.FormData`, but is to contain minimum + information required by TSFC. This class can also combine + multiple `ufl.FormData`s associated with forms that are + extracted form a single form (`original_form`) in one way or + another. This is useful when the user wants to insert some + form specific operations after compiling ufl and before + compiling gem: + + split compile operations compile + UFL gem + op_0 + +--- form_0 ---- gem_0' ---- gem_0 ---+ + | op_1 | + +--- form_1 ---- gem_1' ---- gem_1 ---+ + original_form ---| |---> Kernel + | : : : | + | op_N | + +--- form_N ---- gem_N' ---- gem_N ---+ + + After preprocessing `ufl.FormData`s here: + * Only essential information about the `ufl.FormData`s is retained. + * TSFC can forget `ufl.FormData.original_form`, + * `KernelBuilder`s only need to deal with raw `ufl.Coefficient`s. + + Illustration of the structures. + _____________TSFCFormData____________ + ____________________ __________ | ________ ________ ________ | + |Integral||Integral| |Integral| || || | | || + FormData 0 | Data || Data | ... | Data | || || | | || + |____0___||____1___|_ _|____M___| || TSFC || TSFC | | TSFC || + ____________________ __________ ||Integral||Integral| |Integral|| + |Integral||Integral| |Integral| || Data || Data | | Data || + FormData 1 | Data || Data | ... | Data | || 0 || 1 | | M || + |____0___||____1___|_ _|____M___| ---> || || | ... | || + | | + : : | : : : | + ____________________ __________ | | + |Integral||Integral| |Integral| || || | | || + FormData N | Data || Data | ... | Data | || || | | || + |____0___||____1___|_ _|____M___| ||________||________| |________|| + |_____________________________________| + """ + def __init__(self, form_data_tuple, extraarg_tuple, function_tuple, original_form, diagonal): + arguments = set() + for fd, extraarg in zip(form_data_tuple, extraarg_tuple): + args = [] + for arg in fd.preprocessed_form.arguments(): + if arg not in extraarg: + args.append(arg) + arguments.update((tuple(args), )) + if len(arguments) != 1: + raise ValueError("Found inconsistent sets of arguments in `FormData`s.") + self.arguments, = tuple(arguments) + # Gather all coefficients. + # If a form contains extra arguments, those will be replaced by corresponding functions + # after compiling UFL, so these functions must be included here, too. + reduced_coefficients_set = set(c for fd in form_data_tuple for c in fd.reduced_coefficients) + reduced_coefficients_set.update(chain(*function_tuple)) + reduced_coefficients = sorted(reduced_coefficients_set, key=lambda c: c.count()) + # Reconstruct `ufl.Coefficinet`s with count starting at 0. + function_replace_map = {} + for i, func in enumerate(reduced_coefficients): + for fd in form_data_tuple: + if func in fd.function_replace_map: + coeff = fd.function_replace_map[func] + new_coeff = Coefficient(coeff.ufl_function_space(), count=i) + function_replace_map[func] = new_coeff + break + else: + ufl_function_space = FunctionSpace(func.ufl_domain(), func.ufl_element()) + new_coeff = Coefficient(ufl_function_space, count=i) + function_replace_map[func] = new_coeff + self.reduced_coefficients = reduced_coefficients + self.original_coefficient_positions = [i for i, f in enumerate(original_form.coefficients()) + if f in self.reduced_coefficients] + self.function_replace_map = function_replace_map + + # Translate `ufl.IntegralData`s -> `TSFCIntegralData`. + intg_data_info_dict = {} + for form_data_index, form_data in enumerate(form_data_tuple): + for intg_data in form_data.integral_data: + domain = intg_data.domain + integral_type = intg_data.integral_type + subdomain_id = intg_data.subdomain_id + key = (domain, integral_type, subdomain_id) + # Add (intg_data, form_data, form_data_index). + intg_data_info_dict.setdefault(key, []).append((intg_data, form_data, form_data_index)) + integral_data_list = [] + for key, intg_data_info in intg_data_info_dict.items(): + domain, _, _ = key + domain_number = original_form.domain_numbering()[domain] + integral_data_list.append(TSFCIntegralData(key, intg_data_info, + self, domain_number, function_tuple)) + self.integral_data = tuple(integral_data_list) + + +class TSFCIntegralData(object): + r"""Mimics `ufl.IntegralData`. + + :arg integral_data_key: (domain, integral_type, subdomain_id) tuple. + :arg integral_data_info: A tuple of the lists of integral_data, + form_data, and form_data_index. + :arg tsfc_form_data: The `TSFCFormData` that is to contain this + `TSFCIntegralData` object. + :arg domain_number: The domain number associated with `domain`. + :arg function_tuple: A tuple of functions. + + After preprocessing `ufl.IntegralData`s here: + * Only essential information about the `ufl.IntegralData`s is retained. + * TSFC can forget `ufl.IntegralData.enabled_coefficients`, + """ + def __init__(self, integral_data_key, intg_data_info, tsfc_form_data, domain_number, function_tuple): + self.domain, self.integral_type, self.subdomain_id = integral_data_key + self.domain_number = domain_number + # Gather/preprocess integrals. + integrals = [] + _integral_index_to_form_data_index = [] + functions = set() + for intg_data, form_data, form_data_index in intg_data_info: + for integral in intg_data.integrals: + integrand = integral.integrand() + # Replace functions with Coefficients here. + integrand = ufl.replace(integrand, tsfc_form_data.function_replace_map) + new_integral = integral.reconstruct(integrand=integrand) + integrals.append(new_integral) + # Remember which form_data this integral is associated with. + _integral_index_to_form_data_index.append(form_data_index) + # Gather functions that are enabled in this `TSFCIntegralData`. + functions.update(f for f, enabled in zip(form_data.reduced_coefficients, intg_data.enabled_coefficients) if enabled) + functions.update(function_tuple[form_data_index]) + self.integrals = tuple(integrals) + self._integral_index_to_form_data_index = _integral_index_to_form_data_index + self.arguments = tsfc_form_data.arguments + # This is which coefficient in the original form the + # current coefficient is. + # Ex: + # original_form.coefficients() : f0, f1, f2, f3, f4, f5 + # tsfc_form_data.reduced_coefficients: f1, f2, f3, f5 + # functions : f1, f5 + # self.coefficients : c1, c5 + # self.coefficent_numbers : 1, 5 + functions = sorted(functions, key=lambda c: c.count()) + self.coefficients = tuple(tsfc_form_data.function_replace_map[f] for f in functions) + self.coefficient_numbers = tuple(tsfc_form_data.original_coefficient_positions[tsfc_form_data.reduced_coefficients.index(f)] for f in functions) + + def integral_index_to_form_data_index(self, integral_index): + r"""Return the form data index given an integral index.""" + return self._integral_index_to_form_data_index[integral_index] + + def compile_form(form, prefix="form", parameters=None, interface=None, coffee=True, diagonal=False): """Compiles a UFL form into a set of assembly kernels. @@ -49,15 +210,22 @@ def compile_form(form, prefix="form", parameters=None, interface=None, coffee=Tr assert isinstance(form, Form) + parameters = preprocess_parameters(parameters) + # Determine whether in complex mode: complex_mode = parameters and is_complex(parameters.get("scalar_type")) - fd = ufl_utils.compute_form_data(form, complex_mode=complex_mode) + + form_data = ufl_utils.compute_form_data(form, complex_mode=complex_mode) + if interface: + interface = partial(interface, function_replace_map=form_data.function_replace_map) + tsfc_form_data = TSFCFormData((form_data, ), ((), ), ((), ), form_data.original_form, diagonal) + logger.info(GREEN % "compute_form_data finished in %g seconds.", time.time() - cpu_time) kernels = [] - for integral_data in fd.integral_data: + for tsfc_integral_data in tsfc_form_data.integral_data: start = time.time() - kernel = compile_integral(integral_data, fd, prefix, parameters, interface=interface, coffee=coffee, diagonal=diagonal) + kernel = compile_integral(tsfc_integral_data, prefix, parameters, interface=interface, coffee=coffee, diagonal=diagonal) if kernel is not None: kernels.append(kernel) logger.info(GREEN % "compile_integral finished in %g seconds.", time.time() - start) @@ -66,23 +234,16 @@ def compile_form(form, prefix="form", parameters=None, interface=None, coffee=Tr return kernels -def compile_integral(integral_data, form_data, prefix, parameters, interface, coffee, *, diagonal=False): +def compile_integral(tsfc_integral_data, prefix, parameters, interface, coffee, *, diagonal=False): """Compiles a UFL integral into an assembly kernel. - :arg integral_data: UFL integral data - :arg form_data: UFL form data + :arg tsfc_integral_data: TSFCIntegralData :arg prefix: kernel name will start with this string :arg parameters: parameters object :arg interface: backend module for the kernel interface :arg diagonal: Are we building a kernel for the diagonal of a rank-2 element tensor? :returns: a kernel constructed by the kernel interface """ - if parameters is None: - parameters = default_parameters() - else: - _ = default_parameters() - _.update(parameters) - parameters = _ if interface is None: if coffee: import tsfc.kernel_interface.firedrake as firedrake_interface_coffee @@ -91,180 +252,37 @@ def compile_integral(integral_data, form_data, prefix, parameters, interface, co # Delayed import, loopy is a runtime dependency import tsfc.kernel_interface.firedrake_loopy as firedrake_interface_loopy interface = firedrake_interface_loopy.KernelBuilder - if coffee: - scalar_type = parameters["scalar_type_c"] - else: - scalar_type = parameters["scalar_type"] - # Remove these here, they're handled below. - if parameters.get("quadrature_degree") in ["auto", "default", None, -1, "-1"]: - del parameters["quadrature_degree"] - if parameters.get("quadrature_rule") in ["auto", "default", None]: - del parameters["quadrature_rule"] - - integral_type = integral_data.integral_type - interior_facet = integral_type.startswith("interior_facet") - mesh = integral_data.domain - cell = integral_data.domain.ufl_cell() - arguments = form_data.preprocessed_form.arguments() - kernel_name = "%s_%s_integral_%s" % (prefix, integral_type, integral_data.subdomain_id) - # Handle negative subdomain_id - kernel_name = kernel_name.replace("-", "_") - - fiat_cell = as_fiat_cell(cell) - integration_dim, entity_ids = lower_integral_type(fiat_cell, integral_type) - - quadrature_indices = [] - - # Dict mapping domains to index in original_form.ufl_domains() - domain_numbering = form_data.original_form.domain_numbering() - builder = interface(integral_type, integral_data.subdomain_id, - domain_numbering[integral_data.domain], - scalar_type, + builder = interface(tsfc_integral_data, + parameters["scalar_type_c"] if coffee else parameters["scalar_type"], + parameters["scalar_type"], diagonal=diagonal) - argument_multiindices = tuple(builder.create_element(arg.ufl_element()).get_indices() - for arg in arguments) - if diagonal: - # Error checking occurs in the builder constructor. - # Diagonal assembly is obtained by using the test indices for - # the trial space as well. - a, _ = argument_multiindices - argument_multiindices = (a, a) - - return_variables = builder.set_arguments(arguments, argument_multiindices) - - builder.set_coordinates(mesh) - builder.set_cell_sizes(mesh) - - builder.set_coefficients(integral_data, form_data) - - # Map from UFL FiniteElement objects to multiindices. This is - # so we reuse Index instances when evaluating the same coefficient - # multiple times with the same table. - # - # We also use the same dict for the unconcatenate index cache, - # which maps index objects to tuples of multiindices. These two - # caches shall never conflict as their keys have different types - # (UFL finite elements vs. GEM index objects). - index_cache = {} - - kernel_cfg = dict(interface=builder, - ufl_cell=cell, - integral_type=integral_type, - integration_dim=integration_dim, - entity_ids=entity_ids, - argument_multiindices=argument_multiindices, - index_cache=index_cache, - scalar_type=parameters["scalar_type"]) - - mode_irs = collections.OrderedDict() - for integral in integral_data.integrals: + # Compile UFL -> gem + for integral in tsfc_integral_data.integrals: params = parameters.copy() params.update(integral.metadata()) # integral metadata overrides - if params.get("quadrature_rule") == "default": - del params["quadrature_rule"] - - mode = pick_mode(params["mode"]) - mode_irs.setdefault(mode, collections.OrderedDict()) - - integrand = ufl.replace(integral.integrand(), form_data.function_replace_map) - integrand = ufl_utils.split_coefficients(integrand, builder.coefficient_split) - - # Check if the integral has a quad degree attached, otherwise use - # the estimated polynomial degree attached by compute_form_data - quadrature_degree = params.get("quadrature_degree", - params["estimated_polynomial_degree"]) - try: - quadrature_degree = params["quadrature_degree"] - except KeyError: - quadrature_degree = params["estimated_polynomial_degree"] - functions = list(arguments) + [builder.coordinate(mesh)] + list(integral_data.integral_coefficients) - function_degrees = [f.ufl_function_space().ufl_element().degree() for f in functions] - if all((asarray(quadrature_degree) > 10 * asarray(degree)).all() - for degree in function_degrees): - logger.warning("Estimated quadrature degree %s more " - "than tenfold greater than any " - "argument/coefficient degree (max %s)", - quadrature_degree, max_degree(function_degrees)) - - try: - quad_rule = params["quadrature_rule"] - except KeyError: - integration_cell = fiat_cell.construct_subelement(integration_dim) - quad_rule = make_quadrature(integration_cell, quadrature_degree) - - if not isinstance(quad_rule, AbstractQuadratureRule): - raise ValueError("Expected to find a QuadratureRule object, not a %s" % - type(quad_rule)) - - quadrature_multiindex = quad_rule.point_set.indices - quadrature_indices.extend(quadrature_multiindex) + integrand_exprs = builder.compile_ufl(integral.integrand(), params) + integral_exprs = builder.construct_integrals(integrand_exprs, params) + builder.stash_integrals(integral_exprs, params) + # Compile gem -> kernel + kernel_name = "%s_%s_integral_%s" % (prefix, tsfc_integral_data.integral_type, tsfc_integral_data.subdomain_id) + kernel_name = kernel_name.replace("-", "_") # Handle negative subdomain_id + return builder.construct_kernel(kernel_name) - config = kernel_cfg.copy() - config.update(quadrature_rule=quad_rule) - expressions = fem.compile_ufl(integrand, - interior_facet=interior_facet, - **config) - reps = mode.Integrals(expressions, quadrature_multiindex, - argument_multiindices, params) - for var, rep in zip(return_variables, reps): - mode_irs[mode].setdefault(var, []).append(rep) - - # Finalise mode representations into a set of assignments - assignments = [] - for mode, var_reps in mode_irs.items(): - assignments.extend(mode.flatten(var_reps.items(), index_cache)) - - if assignments: - return_variables, expressions = zip(*assignments) - else: - return_variables = [] - expressions = [] - - # Need optimised roots - options = dict(reduce(operator.and_, - [mode.finalise_options.items() - for mode in mode_irs.keys()])) - expressions = impero_utils.preprocess_gem(expressions, **options) - assignments = list(zip(return_variables, expressions)) - - # Let the kernel interface inspect the optimised IR to register - # what kind of external data is required (e.g., cell orientations, - # cell sizes, etc.). - builder.register_requirements(expressions) - - # Construct ImperoC - split_argument_indices = tuple(chain(*[var.index_ordering() - for var in return_variables])) - index_ordering = tuple(quadrature_indices) + split_argument_indices - try: - impero_c = impero_utils.compile_gem(assignments, index_ordering, remove_zeros=True) - except impero_utils.NoopError: - # No operations, construct empty kernel - return builder.construct_empty_kernel(kernel_name) - - # Generate COFFEE - index_names = [] - - def name_index(index, name): - index_names.append((index, name)) - if index in index_cache: - for multiindex, suffix in zip(index_cache[index], - string.ascii_lowercase): - name_multiindex(multiindex, name + suffix) - - def name_multiindex(multiindex, name): - if len(multiindex) == 1: - name_index(multiindex[0], name) - else: - for i, index in enumerate(multiindex): - name_index(index, name + str(i)) - name_multiindex(quadrature_indices, 'ip') - for multiindex, name in zip(argument_multiindices, ['j', 'k']): - name_multiindex(multiindex, name) - - return builder.construct_kernel(kernel_name, impero_c, index_names, quad_rule) +def preprocess_parameters(parameters): + if parameters is None: + parameters = default_parameters() + else: + _ = default_parameters() + _.update(parameters) + parameters = _ + # Remove these here, they're handled later on. + if parameters.get("quadrature_degree") in ["auto", "default", None, -1, "-1"]: + del parameters["quadrature_degree"] + if parameters.get("quadrature_rule") in ["auto", "default", None]: + del parameters["quadrature_rule"] + return parameters def compile_expression_dual_evaluation(expression, to_element, *, @@ -345,7 +363,7 @@ def compile_expression_dual_evaluation(expression, to_element, *, builder.set_coefficients(coefficients) # Split mixed coefficients - expression = ufl_utils.split_coefficients(expression, builder.coefficient_split) + expression = ufl_utils.split_coefficients(expression, builder.coefficient_split, ) # Translate to GEM kernel_cfg = dict(interface=builder, @@ -435,71 +453,3 @@ def compile_expression_dual_evaluation(expression, to_element, *, builder.register_requirements([ir]) # Build kernel tuple return builder.construct_kernel(return_arg, impero_c, index_names, first_coefficient_fake_coords) - - -def lower_integral_type(fiat_cell, integral_type): - """Lower integral type into the dimension of the integration - subentity and a list of entity numbers for that dimension. - - :arg fiat_cell: FIAT reference cell - :arg integral_type: integral type (string) - """ - vert_facet_types = ['exterior_facet_vert', 'interior_facet_vert'] - horiz_facet_types = ['exterior_facet_bottom', 'exterior_facet_top', 'interior_facet_horiz'] - - dim = fiat_cell.get_dimension() - if integral_type == 'cell': - integration_dim = dim - elif integral_type in ['exterior_facet', 'interior_facet']: - if isinstance(fiat_cell, TensorProductCell): - raise ValueError("{} integral cannot be used with a TensorProductCell; need to distinguish between vertical and horizontal contributions.".format(integral_type)) - integration_dim = dim - 1 - elif integral_type == 'vertex': - integration_dim = 0 - elif integral_type in vert_facet_types + horiz_facet_types: - # Extrusion case - if not isinstance(fiat_cell, TensorProductCell): - raise ValueError("{} integral requires a TensorProductCell.".format(integral_type)) - basedim, extrdim = dim - assert extrdim == 1 - - if integral_type in vert_facet_types: - integration_dim = (basedim - 1, 1) - elif integral_type in horiz_facet_types: - integration_dim = (basedim, 0) - else: - raise NotImplementedError("integral type %s not supported" % integral_type) - - if integral_type == 'exterior_facet_bottom': - entity_ids = [0] - elif integral_type == 'exterior_facet_top': - entity_ids = [1] - else: - entity_ids = list(range(len(fiat_cell.get_topology()[integration_dim]))) - - return integration_dim, entity_ids - - -def pick_mode(mode): - "Return one of the specialized optimisation modules from a mode string." - try: - from firedrake_citations import Citations - cites = {"vanilla": ("Homolya2017", ), - "coffee": ("Luporini2016", "Homolya2017", ), - "spectral": ("Luporini2016", "Homolya2017", "Homolya2017a"), - "tensor": ("Kirby2006", "Homolya2017", )} - for c in cites[mode]: - Citations().register(c) - except ImportError: - pass - if mode == "vanilla": - import tsfc.vanilla as m - elif mode == "coffee": - import tsfc.coffee_mode as m - elif mode == "spectral": - import tsfc.spectral as m - elif mode == "tensor": - import tsfc.tensor as m - else: - raise ValueError("Unknown mode: {}".format(mode)) - return m diff --git a/tsfc/kernel_interface/common.py b/tsfc/kernel_interface/common.py index 9b78817c..2ab7fe76 100644 --- a/tsfc/kernel_interface/common.py +++ b/tsfc/kernel_interface/common.py @@ -1,12 +1,29 @@ +import collections +import string +import operator +from functools import reduce +from itertools import chain + import numpy +from numpy import asarray import coffee.base as coffee +from ufl.utils.sequences import max_degree + import gem from gem.utils import cached_property +import gem.impero_utils as impero_utils +from tsfc import fem, ufl_utils from tsfc.kernel_interface import KernelInterface +from tsfc.finatinterface import as_fiat_cell +from tsfc.logging import logger + +from FIAT.reference_element import TensorProductCell + +from finat.quadrature import AbstractQuadratureRule, make_quadrature class KernelBuilderBase(KernelInterface): @@ -107,3 +124,246 @@ def register_requirements(self, ir): """ # Nothing is required by default pass + + +class KernelBuilderMixin(object): + """Mixin for KernelBuilder classes.""" + + def compile_ufl(self, integrand, params, argument_multiindices=None): + """Compile UFL integrand. + + :arg integrand: UFL integrand. + :arg params: a dict of parameters containing quadrature info. + :kwarg argument_multiindices: multiindices to use to index + arguments contained in the given integrand. If None, the + "true" argument multiindices (`self.argument_multiindices`) + used in the `self.return_variables` are used. + + .. note:: + Problem solving environments can pass any multiindices to be + used in the returned gem expression, e.g., + `self.argument_multiindices_dummy`. They are then responsible + for applying appropriate operations to the returned gem + expression so that the indices in the resulting gem expression + match those in `self.return_variables`. + """ + # Split Coefficients + if self.coefficient_split: + integrand = ufl_utils.split_coefficients(integrand, self.coefficient_split) + # Compile: ufl -> gem + functions = list(self.arguments) + [self.coordinate(self.integral_data.domain)] + list(self.integral_data.coefficients) + _set_quad_rule(params, self.integral_data.domain.ufl_cell(), self.integral_data.integral_type, functions) + quad_rule = params["quadrature_rule"] + config = self.fem_config.copy() + config.update(quadrature_rule=quad_rule) + config['argument_multiindices'] = argument_multiindices or self.argument_multiindices + expressions = fem.compile_ufl(integrand, + interior_facet=self.interior_facet, + **config) + self.quadrature_indices.extend(quad_rule.point_set.indices) + return expressions + + def construct_integrals(self, expressions, params): + mode = pick_mode(params["mode"]) + return mode.Integrals(expressions, + params["quadrature_rule"].point_set.indices, + self.argument_multiindices, + params) + + def stash_integrals(self, reps, params): + mode = pick_mode(params["mode"]) + mode_irs = self.mode_irs + mode_irs.setdefault(mode, collections.OrderedDict()) + for var, rep in zip(self.return_variables, reps): + mode_irs[mode].setdefault(var, []).append(rep) + + def compile_gem(self): + # Finalise mode representations into a set of assignments + mode_irs = self.mode_irs + index_cache = self.fem_config['index_cache'] + + assignments = [] + for mode, var_reps in mode_irs.items(): + assignments.extend(mode.flatten(var_reps.items(), index_cache)) + + if assignments: + return_variables, expressions = zip(*assignments) + else: + return_variables = [] + expressions = [] + + # Need optimised roots + options = dict(reduce(operator.and_, + [mode.finalise_options.items() + for mode in mode_irs.keys()])) + expressions = impero_utils.preprocess_gem(expressions, **options) + + # Let the kernel interface inspect the optimised IR to register + # what kind of external data is required (e.g., cell orientations, + # cell sizes, etc.). + oriented, needs_cell_sizes, tabulations = self.register_requirements(expressions) + + # Construct ImperoC + assignments = list(zip(return_variables, expressions)) + index_ordering = get_index_ordering(self.quadrature_indices, return_variables) + try: + impero_c = impero_utils.compile_gem(assignments, index_ordering, remove_zeros=True) + except impero_utils.NoopError: + impero_c = None + return impero_c, oriented, needs_cell_sizes, tabulations + + @cached_property + def argument_multiindices_dummy(self): + return tuple(tuple(gem.Index(extent=a.extent) for a in arg) + for arg in self.argument_multiindices) + + @cached_property + def fem_config(self): + # Map from UFL FiniteElement objects to multiindices. This is + # so we reuse Index instances when evaluating the same coefficient + # multiple times with the same table. + # + # We also use the same dict for the unconcatenate index cache, + # which maps index objects to tuples of multiindices. These two + # caches shall never conflict as their keys have different types + # (UFL finite elements vs. GEM index objects). + # + # -> fem_config['index_cache'] + integral_type = self.integral_data.integral_type + cell = self.integral_data.domain.ufl_cell() + fiat_cell = as_fiat_cell(cell) + integration_dim, entity_ids = lower_integral_type(fiat_cell, integral_type) + return dict(interface=self, + ufl_cell=cell, + integral_type=integral_type, + integration_dim=integration_dim, + entity_ids=entity_ids, + index_cache={}, + scalar_type=self.fem_scalar_type) + + +def get_index_ordering(quadrature_indices, return_variables): + split_argument_indices = tuple(chain(*[var.index_ordering() + for var in return_variables])) + return tuple(quadrature_indices) + split_argument_indices + + +def get_index_names(quadrature_indices, argument_multiindices, index_cache): + index_names = [] + + def name_index(index, name): + index_names.append((index, name)) + if index in index_cache: + for multiindex, suffix in zip(index_cache[index], + string.ascii_lowercase): + name_multiindex(multiindex, name + suffix) + + def name_multiindex(multiindex, name): + if len(multiindex) == 1: + name_index(multiindex[0], name) + else: + for i, index in enumerate(multiindex): + name_index(index, name + str(i)) + + name_multiindex(quadrature_indices, 'ip') + for multiindex, name in zip(argument_multiindices, ['j', 'k']): + name_multiindex(multiindex, name) + return index_names + + +def _set_quad_rule(params, cell, integral_type, functions): + # Check if the integral has a quad degree attached, otherwise use + # the estimated polynomial degree attached by compute_form_data + try: + quadrature_degree = params["quadrature_degree"] + except KeyError: + quadrature_degree = params["estimated_polynomial_degree"] + function_degrees = [f.ufl_function_space().ufl_element().degree() for f in functions] + if all((asarray(quadrature_degree) > 10 * asarray(degree)).all() + for degree in function_degrees): + logger.warning("Estimated quadrature degree %s more " + "than tenfold greater than any " + "argument/coefficient degree (max %s)", + quadrature_degree, max_degree(function_degrees)) + if params.get("quadrature_rule") == "default": + del params["quadrature_rule"] + try: + quad_rule = params["quadrature_rule"] + except KeyError: + fiat_cell = as_fiat_cell(cell) + integration_dim, _ = lower_integral_type(fiat_cell, integral_type) + integration_cell = fiat_cell.construct_subelement(integration_dim) + quad_rule = make_quadrature(integration_cell, quadrature_degree) + params["quadrature_rule"] = quad_rule + + if not isinstance(quad_rule, AbstractQuadratureRule): + raise ValueError("Expected to find a QuadratureRule object, not a %s" % + type(quad_rule)) + + +def lower_integral_type(fiat_cell, integral_type): + """Lower integral type into the dimension of the integration + subentity and a list of entity numbers for that dimension. + + :arg fiat_cell: FIAT reference cell + :arg integral_type: integral type (string) + """ + vert_facet_types = ['exterior_facet_vert', 'interior_facet_vert'] + horiz_facet_types = ['exterior_facet_bottom', 'exterior_facet_top', 'interior_facet_horiz'] + + dim = fiat_cell.get_dimension() + if integral_type == 'cell': + integration_dim = dim + elif integral_type in ['exterior_facet', 'interior_facet']: + if isinstance(fiat_cell, TensorProductCell): + raise ValueError("{} integral cannot be used with a TensorProductCell; need to distinguish between vertical and horizontal contributions.".format(integral_type)) + integration_dim = dim - 1 + elif integral_type == 'vertex': + integration_dim = 0 + elif integral_type in vert_facet_types + horiz_facet_types: + # Extrusion case + if not isinstance(fiat_cell, TensorProductCell): + raise ValueError("{} integral requires a TensorProductCell.".format(integral_type)) + basedim, extrdim = dim + assert extrdim == 1 + + if integral_type in vert_facet_types: + integration_dim = (basedim - 1, 1) + elif integral_type in horiz_facet_types: + integration_dim = (basedim, 0) + else: + raise NotImplementedError("integral type %s not supported" % integral_type) + + if integral_type == 'exterior_facet_bottom': + entity_ids = [0] + elif integral_type == 'exterior_facet_top': + entity_ids = [1] + else: + entity_ids = list(range(len(fiat_cell.get_topology()[integration_dim]))) + + return integration_dim, entity_ids + + +def pick_mode(mode): + "Return one of the specialized optimisation modules from a mode string." + try: + from firedrake_citations import Citations + cites = {"vanilla": ("Homolya2017", ), + "coffee": ("Luporini2016", "Homolya2017", ), + "spectral": ("Luporini2016", "Homolya2017", "Homolya2017a"), + "tensor": ("Kirby2006", "Homolya2017", )} + for c in cites[mode]: + Citations().register(c) + except ImportError: + pass + if mode == "vanilla": + import tsfc.vanilla as m + elif mode == "coffee": + import tsfc.coffee_mode as m + elif mode == "spectral": + import tsfc.spectral as m + elif mode == "tensor": + import tsfc.tensor as m + else: + raise ValueError("Unknown mode: {}".format(mode)) + return m diff --git a/tsfc/kernel_interface/firedrake.py b/tsfc/kernel_interface/firedrake.py index 5c5c3ee5..8f159177 100644 --- a/tsfc/kernel_interface/firedrake.py +++ b/tsfc/kernel_interface/firedrake.py @@ -1,4 +1,5 @@ import numpy +import collections from collections import namedtuple from itertools import chain, product from functools import partial @@ -12,7 +13,7 @@ from gem.optimise import remove_componenttensors as prune from tsfc.finatinterface import create_element -from tsfc.kernel_interface.common import KernelBuilderBase as _KernelBuilderBase +from tsfc.kernel_interface.common import KernelBuilderBase as _KernelBuilderBase, KernelBuilderMixin, get_index_names from tsfc.coffee import generate as generate_coffee @@ -26,8 +27,8 @@ def make_builder(*args, **kwargs): class Kernel(object): __slots__ = ("ast", "integral_type", "oriented", "subdomain_id", - "domain_number", "needs_cell_sizes", "tabulations", "quadrature_rule", - "coefficient_numbers", "__weakref__") + "domain_number", "needs_cell_sizes", "tabulations", + "coefficient_numbers", "external_data_numbers", "external_data_parts", "__weakref__") """A compiled Kernel object. :kwarg ast: The COFFEE ast for the kernel. @@ -39,14 +40,22 @@ class Kernel(object): original_form.ufl_domains() to get the correct domain). :kwarg coefficient_numbers: A list of which coefficients from the form the kernel needs. - :kwarg quadrature_rule: The finat quadrature rule used to generate this kernel + :kwarg external_data_numbers: A list of external data structures + the kernel needs. These data structures do not originate in + UFL forms, but in the operations potentially applied to gem + expressions after compiling UFL and before compiling gem. + :kwarg external_data_parts: A list of tuples of indices. Each + tuple contains indices of the associated external data + structure that the kernel needs. :kwarg tabulations: The runtime tabulations this kernel requires :kwarg needs_cell_sizes: Does the kernel require cell sizes. """ def __init__(self, ast=None, integral_type=None, oriented=False, - subdomain_id=None, domain_number=None, quadrature_rule=None, + subdomain_id=None, domain_number=None, coefficient_numbers=(), - needs_cell_sizes=False): + external_data_numbers=(), external_data_parts=(), + needs_cell_sizes=False, + tabulations=None): # Defaults self.ast = ast self.integral_type = integral_type @@ -54,7 +63,10 @@ def __init__(self, ast=None, integral_type=None, oriented=False, self.domain_number = domain_number self.subdomain_id = subdomain_id self.coefficient_numbers = coefficient_numbers + self.external_data_numbers = external_data_numbers + self.external_data_parts = external_data_parts self.needs_cell_sizes = needs_cell_sizes + self.tabulations = tabulations super(Kernel, self).__init__() @@ -85,7 +97,7 @@ def _coefficient(self, coefficient, name): :arg name: coefficient name :returns: COFFEE function argument for the coefficient """ - funarg, expression = prepare_coefficient(coefficient, name, self.scalar_type, interior_facet=self.interior_facet) + funarg, expression = prepare_coefficient(coefficient.ufl_element(), name, self.scalar_type, interior_facet=self.interior_facet) self.coefficient_map[coefficient] = expression return funarg @@ -107,7 +119,7 @@ def set_cell_sizes(self, domain): # topological_dimension is 0 and the concept of "cell size" # is not useful for a vertex. f = Coefficient(FunctionSpace(domain, FiniteElement("P", domain.ufl_cell(), 1))) - funarg, expression = prepare_coefficient(f, "cell_sizes", self.scalar_type, interior_facet=self.interior_facet) + funarg, expression = prepare_coefficient(f.ufl_element(), "cell_sizes", self.scalar_type, interior_facet=self.interior_facet) self.cell_sizes_arg = funarg self._cell_sizes = expression @@ -174,24 +186,25 @@ def construct_kernel(self, return_arg, impero_c, index_names): return ExpressionKernel(kernel_code, self.oriented, self.cell_sizes, self.coefficients, self.tabulations) -class KernelBuilder(KernelBuilderBase): +class KernelBuilder(KernelBuilderBase, KernelBuilderMixin): """Helper class for building a :class:`Kernel` object.""" - def __init__(self, integral_type, subdomain_id, domain_number, scalar_type, - dont_split=(), diagonal=False): + def __init__(self, integral_data, scalar_type, fem_scalar_type, + dont_split=(), function_replace_map={}, diagonal=False): """Initialise a kernel builder.""" - super(KernelBuilder, self).__init__(scalar_type, integral_type.startswith("interior_facet")) + KernelBuilderBase.__init__(self, scalar_type, integral_data.integral_type.startswith("interior_facet")) + self.fem_scalar_type = fem_scalar_type - self.kernel = Kernel(integral_type=integral_type, subdomain_id=subdomain_id, - domain_number=domain_number) self.diagonal = diagonal - self.local_tensor = None self.coordinates_arg = None self.coefficient_args = [] self.coefficient_split = {} - self.dont_split = frozenset(dont_split) + self.external_data_args = [] + # Map functions to raw coefficients here. + self.dont_split = frozenset(function_replace_map[f] for f in dont_split if f in function_replace_map) # Facet number + integral_type = integral_data.integral_type if integral_type in ['exterior_facet', 'exterior_facet_vert']: facet = gem.Variable('facet', (1,)) self._entity_number = {None: gem.VariableIndex(gem.Indexed(facet, (0,)))} @@ -204,17 +217,38 @@ def __init__(self, integral_type, subdomain_id, domain_number, scalar_type, elif integral_type == 'interior_facet_horiz': self._entity_number = {'+': 1, '-': 0} - def set_arguments(self, arguments, multiindices): + self.set_coordinates(integral_data.domain) + self.set_cell_sizes(integral_data.domain) + self.set_coefficients(integral_data.coefficients) + + self.integral_data = integral_data + self.arguments = integral_data.arguments + self.local_tensor, self.return_variables, self.argument_multiindices = self.set_arguments(self.arguments) + self.mode_irs = collections.OrderedDict() + self.quadrature_indices = [] + + def set_arguments(self, arguments): """Process arguments. :arg arguments: :class:`ufl.Argument`s - :arg multiindices: GEM argument multiindices - :returns: GEM expression representing the return variable + :returns: :class:`loopy.GlobalArg` for the return variable, + GEM expression representing the return variable, + GEM argument multiindices. """ - self.local_tensor, expressions = prepare_arguments( - arguments, multiindices, self.scalar_type, interior_facet=self.interior_facet, - diagonal=self.diagonal) - return expressions + argument_multiindices = tuple(create_element(arg.ufl_element()).get_indices() + for arg in arguments) + if self.diagonal: + # Error checking occurs in the builder constructor. + # Diagonal assembly is obtained by using the test indices for + # the trial space as well. + a, _ = argument_multiindices + argument_multiindices = (a, a) + local_tensor, return_variables = prepare_arguments(arguments, + argument_multiindices, + self.scalar_type, + interior_facet=self.interior_facet, + diagonal=self.diagonal) + return local_tensor, return_variables, argument_multiindices def set_coordinates(self, domain): """Prepare the coordinate field. @@ -226,87 +260,103 @@ def set_coordinates(self, domain): self.domain_coordinate[domain] = f self.coordinates_arg = self._coefficient(f, "coords") - def set_coefficients(self, integral_data, form_data): + def set_coefficients(self, coefficients): """Prepare the coefficients of the form. - :arg integral_data: UFL integral data - :arg form_data: UFL form data + :arg coefficients: a tuple of `ufl.Coefficient`s. """ - coefficients = [] - coefficient_numbers = [] - # enabled_coefficients is a boolean array that indicates which - # of reduced_coefficients the integral requires. - for i in range(len(integral_data.enabled_coefficients)): - if integral_data.enabled_coefficients[i]: - original = form_data.reduced_coefficients[i] - coefficient = form_data.function_replace_map[original] - if type(coefficient.ufl_element()) == ufl_MixedElement: - if original in self.dont_split: - coefficients.append(coefficient) - self.coefficient_split[coefficient] = [coefficient] - else: - split = [Coefficient(FunctionSpace(coefficient.ufl_domain(), element)) - for element in coefficient.ufl_element().sub_elements()] - coefficients.extend(split) - self.coefficient_split[coefficient] = split + coeffs = [] + for c in coefficients: + if type(c.ufl_element()) == ufl_MixedElement: + if c in self.dont_split: + coeffs.append(c) + self.coefficient_split[c] = [c] else: - coefficients.append(coefficient) - # This is which coefficient in the original form the - # current coefficient is. - # Consider f*v*dx + g*v*ds, the full form contains two - # coefficients, but each integral only requires one. - coefficient_numbers.append(form_data.original_coefficient_positions[i]) - for i, coefficient in enumerate(coefficients): - self.coefficient_args.append( - self._coefficient(coefficient, "w_%d" % i)) - self.kernel.coefficient_numbers = tuple(coefficient_numbers) + split = [Coefficient(FunctionSpace(c.ufl_domain(), element)) + for element in c.ufl_element().sub_elements()] + coeffs.extend(split) + self.coefficient_split[c] = split + else: + coeffs.append(c) + for i, c in enumerate(coeffs): + self.coefficient_args.append(self._coefficient(c, "w_%d" % i)) + + def set_external_data(self, elements): + """Prepare external data structures. + + :arg elements: a tuple of `ufl.FiniteElement`s. + :returns: gem expressions for the data represented by elements. + + The retuned gem expressions are to be used in the operations + applied to the gem expressions obtained by compiling UFL before + compiling gem. The users are responsible for bridging these + gem expressions and actual data by setting correct values in + `external_data_numbers` and `external_data_parts` in the kernel. + """ + if any(type(element) == ufl_MixedElement for element in elements): + raise ValueError("Unable to handle `MixedElement`s.") + expressions = [] + for i, element in enumerate(elements): + funarg, expression = prepare_coefficient(element, "e_%d" % i, self.scalar_type, interior_facet=self.interior_facet) + self.external_data_args.append(funarg) + expressions.append(expression) + return tuple(expressions) def register_requirements(self, ir): """Inspect what is referenced by the IR that needs to be provided by the kernel interface.""" - knl = self.kernel - knl.oriented, knl.needs_cell_sizes, knl.tabulations = check_requirements(ir) + return check_requirements(ir) - def construct_kernel(self, name, impero_c, index_names, quadrature_rule): + def construct_kernel(self, kernel_name, external_data_numbers=(), external_data_parts=()): """Construct a fully built :class:`Kernel`. This function contains the logic for building the argument list for assembly kernels. - :arg name: function name - :arg impero_c: ImperoC tuple with Impero AST and other data - :arg index_names: pre-assigned index names - :arg quadrature rule: quadrature rule - + :arg kernel_name: function name :returns: :class:`Kernel` object """ - body = generate_coffee(impero_c, index_names, self.scalar_type) + integral_data = self.integral_data + + impero_c, oriented, needs_cell_sizes, tabulations = self.compile_gem() + + if impero_c is None: + return self.construct_empty_kernel(kernel_name) args = [self.local_tensor, self.coordinates_arg] - if self.kernel.oriented: + if oriented: args.append(cell_orientations_coffee_arg) - if self.kernel.needs_cell_sizes: + if needs_cell_sizes: args.append(self.cell_sizes_arg) - args.extend(self.coefficient_args) - if self.kernel.integral_type in ["exterior_facet", "exterior_facet_vert"]: + args.extend(self.coefficient_args + self.external_data_args) + if integral_data.integral_type in ["exterior_facet", "exterior_facet_vert"]: args.append(coffee.Decl("unsigned int", coffee.Symbol("facet", rank=(1,)), qualifiers=["const"])) - elif self.kernel.integral_type in ["interior_facet", "interior_facet_vert"]: + elif integral_data.integral_type in ["interior_facet", "interior_facet_vert"]: args.append(coffee.Decl("unsigned int", coffee.Symbol("facet", rank=(2,)), qualifiers=["const"])) - - for name_, shape in self.kernel.tabulations: + for name, shape in tabulations: args.append(coffee.Decl(self.scalar_type, coffee.Symbol( - name_, rank=shape), qualifiers=["const"])) - - self.kernel.quadrature_rule = quadrature_rule - - self.kernel.ast = KernelBuilderBase.construct_kernel(self, name, args, body) - return self.kernel - - def construct_empty_kernel(self, name): + name, rank=shape), qualifiers=["const"])) + index_cache = self.fem_config['index_cache'] + index_names = get_index_names(self.quadrature_indices, self.argument_multiindices, index_cache) + body = generate_coffee(impero_c, index_names, self.scalar_type) + ast = KernelBuilderBase.construct_kernel(self, kernel_name, args, body) + + return Kernel(ast=ast, + integral_type=integral_data.integral_type, + subdomain_id=integral_data.subdomain_id, + domain_number=integral_data.domain_number, + coefficient_numbers=integral_data.coefficient_numbers, + external_data_numbers=external_data_numbers, + external_data_parts=external_data_parts, + oriented=oriented, + needs_cell_sizes=needs_cell_sizes, + tabulations=tabulations) + + def construct_empty_kernel(self, kernel_name): """Return None, since Firedrake needs no empty kernels. :arg name: function name @@ -332,11 +382,11 @@ def check_requirements(ir): return cell_orientations, cell_sizes, tuple(sorted(rt_tabs.items())) -def prepare_coefficient(coefficient, name, scalar_type, interior_facet=False): +def prepare_coefficient(ufl_element, name, scalar_type, interior_facet=False): """Bridges the kernel interface and the GEM abstraction for Coefficients. - :arg coefficient: UFL Coefficient + :arg ufl_element: UFL element :arg name: unique name to refer to the Coefficient in the kernel :arg interior_facet: interior facet integral? :returns: (funarg, expression) @@ -346,18 +396,18 @@ def prepare_coefficient(coefficient, name, scalar_type, interior_facet=False): """ assert isinstance(interior_facet, bool) - if coefficient.ufl_element().family() == 'Real': + if ufl_element.family() == 'Real': # Constant funarg = coffee.Decl(scalar_type, coffee.Symbol(name), pointers=[("restrict",)], qualifiers=["const"]) expression = gem.reshape(gem.Variable(name, (None,)), - coefficient.ufl_shape) + ufl_element.value_shape()) return funarg, expression - finat_element = create_element(coefficient.ufl_element()) + finat_element = create_element(ufl_element) shape = finat_element.index_shape size = numpy.prod(shape, dtype=int) diff --git a/tsfc/kernel_interface/firedrake_loopy.py b/tsfc/kernel_interface/firedrake_loopy.py index eb90cf72..e066e568 100644 --- a/tsfc/kernel_interface/firedrake_loopy.py +++ b/tsfc/kernel_interface/firedrake_loopy.py @@ -1,4 +1,5 @@ import numpy +import collections from collections import namedtuple from itertools import chain, product from functools import partial @@ -11,7 +12,7 @@ import loopy as lp from tsfc.finatinterface import create_element -from tsfc.kernel_interface.common import KernelBuilderBase as _KernelBuilderBase +from tsfc.kernel_interface.common import KernelBuilderBase as _KernelBuilderBase, KernelBuilderMixin, get_index_names from tsfc.kernel_interface.firedrake import check_requirements from tsfc.loopy import generate as generate_loopy @@ -26,8 +27,8 @@ def make_builder(*args, **kwargs): class Kernel(object): __slots__ = ("ast", "integral_type", "oriented", "subdomain_id", - "domain_number", "needs_cell_sizes", "tabulations", "quadrature_rule", - "coefficient_numbers", "__weakref__") + "domain_number", "needs_cell_sizes", "tabulations", + "coefficient_numbers", "external_data_numbers", "external_data_parts", "__weakref__") """A compiled Kernel object. :kwarg ast: The loopy kernel object. @@ -39,14 +40,22 @@ class Kernel(object): original_form.ufl_domains() to get the correct domain). :kwarg coefficient_numbers: A list of which coefficients from the form the kernel needs. - :kwarg quadrature_rule: The finat quadrature rule used to generate this kernel + :kwarg external_data_numbers: A list of external data structures + the kernel needs. These data structures do not originate in + UFL forms, but in the operations potentially applied to gem + expressions after compiling UFL and before compiling gem. + :kwarg external_data_parts: A list of tuples of indices. Each + tuple contains indices of the associated external data + structure that the kernel needs. :kwarg tabulations: The runtime tabulations this kernel requires :kwarg needs_cell_sizes: Does the kernel require cell sizes. """ def __init__(self, ast=None, integral_type=None, oriented=False, - subdomain_id=None, domain_number=None, quadrature_rule=None, + subdomain_id=None, domain_number=None, coefficient_numbers=(), - needs_cell_sizes=False): + external_data_numbers=(), external_data_parts=(), + needs_cell_sizes=False, + tabulations=None): # Defaults self.ast = ast self.integral_type = integral_type @@ -54,7 +63,10 @@ def __init__(self, ast=None, integral_type=None, oriented=False, self.domain_number = domain_number self.subdomain_id = subdomain_id self.coefficient_numbers = coefficient_numbers + self.external_data_numbers = external_data_numbers + self.external_data_parts = external_data_parts self.needs_cell_sizes = needs_cell_sizes + self.tabulations = tabulations super(Kernel, self).__init__() @@ -88,7 +100,7 @@ def _coefficient(self, coefficient, name): :arg name: coefficient name :returns: loopy argument for the coefficient """ - funarg, expression = prepare_coefficient(coefficient, name, self.scalar_type, interior_facet=self.interior_facet) + funarg, expression = prepare_coefficient(coefficient.ufl_element(), name, self.scalar_type, interior_facet=self.interior_facet) self.coefficient_map[coefficient] = expression return funarg @@ -110,7 +122,7 @@ def set_cell_sizes(self, domain): # topological_dimension is 0 and the concept of "cell size" # is not useful for a vertex. f = Coefficient(FunctionSpace(domain, FiniteElement("P", domain.ufl_cell(), 1))) - funarg, expression = prepare_coefficient(f, "cell_sizes", self.scalar_type, interior_facet=self.interior_facet) + funarg, expression = prepare_coefficient(f.ufl_element(), "cell_sizes", self.scalar_type, interior_facet=self.interior_facet) self.cell_sizes_arg = funarg self._cell_sizes = expression @@ -124,7 +136,7 @@ class ExpressionKernelBuilder(KernelBuilderBase): """Builds expression kernels for UFL interpolation in Firedrake.""" def __init__(self, scalar_type): - super(ExpressionKernelBuilder, self).__init__(scalar_type=scalar_type) + super(ExpressionKernelBuilder, self).__init__(scalar_type) self.oriented = False self.cell_sizes = False @@ -179,24 +191,25 @@ def construct_kernel(self, return_arg, impero_c, index_names, first_coefficient_ self.tabulations) -class KernelBuilder(KernelBuilderBase): +class KernelBuilder(KernelBuilderBase, KernelBuilderMixin): """Helper class for building a :class:`Kernel` object.""" - def __init__(self, integral_type, subdomain_id, domain_number, scalar_type, dont_split=(), - diagonal=False): + def __init__(self, integral_data, scalar_type, fem_scalar_type, + dont_split=(), function_replace_map={}, diagonal=False): """Initialise a kernel builder.""" - super(KernelBuilder, self).__init__(scalar_type, integral_type.startswith("interior_facet")) + KernelBuilderBase.__init__(self, scalar_type, integral_data.integral_type.startswith("interior_facet")) + self.fem_scalar_type = fem_scalar_type - self.kernel = Kernel(integral_type=integral_type, subdomain_id=subdomain_id, - domain_number=domain_number) self.diagonal = diagonal - self.local_tensor = None self.coordinates_arg = None self.coefficient_args = [] self.coefficient_split = {} - self.dont_split = frozenset(dont_split) + self.external_data_args = [] + # Map functions to raw coefficients here. + self.dont_split = frozenset(function_replace_map[f] for f in dont_split if f in function_replace_map) # Facet number + integral_type = integral_data.integral_type if integral_type in ['exterior_facet', 'exterior_facet_vert']: facet = gem.Variable('facet', (1,)) self._entity_number = {None: gem.VariableIndex(gem.Indexed(facet, (0,)))} @@ -209,17 +222,38 @@ def __init__(self, integral_type, subdomain_id, domain_number, scalar_type, dont elif integral_type == 'interior_facet_horiz': self._entity_number = {'+': 1, '-': 0} - def set_arguments(self, arguments, multiindices): + self.set_coordinates(integral_data.domain) + self.set_cell_sizes(integral_data.domain) + self.set_coefficients(integral_data.coefficients) + + self.integral_data = integral_data + self.arguments = integral_data.arguments + self.local_tensor, self.return_variables, self.argument_multiindices = self.set_arguments(self.arguments) + self.mode_irs = collections.OrderedDict() + self.quadrature_indices = [] + + def set_arguments(self, arguments): """Process arguments. :arg arguments: :class:`ufl.Argument`s - :arg multiindices: GEM argument multiindices - :returns: GEM expression representing the return variable + :returns: :class:`loopy.GlobalArg` for the return variable, + GEM expression representing the return variable, + GEM argument multiindices. """ - self.local_tensor, expressions = prepare_arguments( - arguments, multiindices, self.scalar_type, interior_facet=self.interior_facet, - diagonal=self.diagonal) - return expressions + argument_multiindices = tuple(create_element(arg.ufl_element()).get_indices() + for arg in arguments) + if self.diagonal: + # Error checking occurs in the builder constructor. + # Diagonal assembly is obtained by using the test indices for + # the trial space as well. + a, _ = argument_multiindices + argument_multiindices = (a, a) + local_tensor, return_variables = prepare_arguments(arguments, + argument_multiindices, + self.scalar_type, + interior_facet=self.interior_facet, + diagonal=self.diagonal) + return local_tensor, return_variables, argument_multiindices def set_coordinates(self, domain): """Prepare the coordinate field. @@ -231,79 +265,99 @@ def set_coordinates(self, domain): self.domain_coordinate[domain] = f self.coordinates_arg = self._coefficient(f, "coords") - def set_coefficients(self, integral_data, form_data): + def set_coefficients(self, coefficients): """Prepare the coefficients of the form. - :arg integral_data: UFL integral data - :arg form_data: UFL form data + :arg coefficients: a tuple of `ufl.Coefficient`s. """ - coefficients = [] - coefficient_numbers = [] - # enabled_coefficients is a boolean array that indicates which - # of reduced_coefficients the integral requires. - for i in range(len(integral_data.enabled_coefficients)): - if integral_data.enabled_coefficients[i]: - original = form_data.reduced_coefficients[i] - coefficient = form_data.function_replace_map[original] - if type(coefficient.ufl_element()) == ufl_MixedElement: - if original in self.dont_split: - coefficients.append(coefficient) - self.coefficient_split[coefficient] = [coefficient] - else: - split = [Coefficient(FunctionSpace(coefficient.ufl_domain(), element)) - for element in coefficient.ufl_element().sub_elements()] - coefficients.extend(split) - self.coefficient_split[coefficient] = split + coeffs = [] + for c in coefficients: + if type(c.ufl_element()) == ufl_MixedElement: + if c in self.dont_split: + coeffs.append(c) + self.coefficient_split[c] = [c] else: - coefficients.append(coefficient) - # This is which coefficient in the original form the - # current coefficient is. - # Consider f*v*dx + g*v*ds, the full form contains two - # coefficients, but each integral only requires one. - coefficient_numbers.append(form_data.original_coefficient_positions[i]) - for i, coefficient in enumerate(coefficients): - self.coefficient_args.append( - self._coefficient(coefficient, "w_%d" % i)) - self.kernel.coefficient_numbers = tuple(coefficient_numbers) + split = [Coefficient(FunctionSpace(c.ufl_domain(), element)) + for element in c.ufl_element().sub_elements()] + coeffs.extend(split) + self.coefficient_split[c] = split + else: + coeffs.append(c) + for i, c in enumerate(coeffs): + self.coefficient_args.append(self._coefficient(c, "w_%d" % i)) + + def set_external_data(self, elements): + """Prepare external data structures. + + :arg elements: a tuple of `ufl.FiniteElement`s. + :returns: gem expressions for the data represented by elements. + + The retuned gem expressions are to be used in the operations + applied to the gem expressions obtained by compiling UFL before + compiling gem. The users are responsible for bridging these + gem expressions and actual data by setting correct values in + `external_data_numbers` and `external_data_parts` in the kernel. + """ + if any(type(element) == ufl_MixedElement for element in elements): + raise ValueError("Unable to handle `MixedElement`s.") + expressions = [] + for i, element in enumerate(elements): + funarg, expression = prepare_coefficient(element, "e_%d" % i, self.scalar_type, interior_facet=self.interior_facet) + self.external_data_args.append(funarg) + expressions.append(expression) + return tuple(expressions) def register_requirements(self, ir): """Inspect what is referenced by the IR that needs to be provided by the kernel interface.""" - knl = self.kernel - knl.oriented, knl.needs_cell_sizes, knl.tabulations = check_requirements(ir) + return check_requirements(ir) - def construct_kernel(self, name, impero_c, index_names, quadrature_rule): + def construct_kernel(self, kernel_name, external_data_numbers=(), external_data_parts=()): """Construct a fully built :class:`Kernel`. This function contains the logic for building the argument list for assembly kernels. - :arg name: function name - :arg impero_c: ImperoC tuple with Impero AST and other data - :arg index_names: pre-assigned index names - :arg quadrature rule: quadrature rule + :arg kernel_name: function name + :kwarg external_data_numbers: see :class:`Kernel`. + :kwarg external_data_parts: see :class:`Kernel.` :returns: :class:`Kernel` object """ + integral_data = self.integral_data + + impero_c, oriented, needs_cell_sizes, tabulations = self.compile_gem() + + if impero_c is None: + return self.construct_empty_kernel(kernel_name) args = [self.local_tensor, self.coordinates_arg] - if self.kernel.oriented: + if oriented: args.append(self.cell_orientations_loopy_arg) - if self.kernel.needs_cell_sizes: + if needs_cell_sizes: args.append(self.cell_sizes_arg) - args.extend(self.coefficient_args) - if self.kernel.integral_type in ["exterior_facet", "exterior_facet_vert"]: + args.extend(self.coefficient_args + self.external_data_args) + if integral_data.integral_type in ["exterior_facet", "exterior_facet_vert"]: args.append(lp.GlobalArg("facet", dtype=numpy.uint32, shape=(1,))) - elif self.kernel.integral_type in ["interior_facet", "interior_facet_vert"]: + elif integral_data.integral_type in ["interior_facet", "interior_facet_vert"]: args.append(lp.GlobalArg("facet", dtype=numpy.uint32, shape=(2,))) - - for name_, shape in self.kernel.tabulations: - args.append(lp.GlobalArg(name_, dtype=self.scalar_type, shape=shape)) - - self.kernel.quadrature_rule = quadrature_rule - self.kernel.ast = generate_loopy(impero_c, args, self.scalar_type, name, index_names) - return self.kernel - - def construct_empty_kernel(self, name): + for name, shape in tabulations: + args.append(lp.GlobalArg(name, dtype=self.scalar_type, shape=shape)) + index_cache = self.fem_config['index_cache'] + index_names = get_index_names(self.quadrature_indices, self.argument_multiindices, index_cache) + ast = generate_loopy(impero_c, args, self.scalar_type, kernel_name, index_names) + + return Kernel(ast=ast, + integral_type=integral_data.integral_type, + subdomain_id=integral_data.subdomain_id, + domain_number=integral_data.domain_number, + coefficient_numbers=integral_data.coefficient_numbers, + external_data_numbers=external_data_numbers, + external_data_parts=external_data_parts, + oriented=oriented, + needs_cell_sizes=needs_cell_sizes, + tabulations=tabulations) + + def construct_empty_kernel(self, kernel_name): """Return None, since Firedrake needs no empty kernels. :arg name: function name @@ -312,11 +366,11 @@ def construct_empty_kernel(self, name): return None -def prepare_coefficient(coefficient, name, scalar_type, interior_facet=False): +def prepare_coefficient(ufl_element, name, scalar_type, interior_facet=False): """Bridges the kernel interface and the GEM abstraction for Coefficients. - :arg coefficient: UFL Coefficient + :arg ufl_element: UFL element :arg name: unique name to refer to the Coefficient in the kernel :arg interior_facet: interior facet integral? :returns: (funarg, expression) @@ -326,25 +380,24 @@ def prepare_coefficient(coefficient, name, scalar_type, interior_facet=False): """ assert isinstance(interior_facet, bool) - if coefficient.ufl_element().family() == 'Real': + if ufl_element.family() == 'Real': # Constant - funarg = lp.GlobalArg(name, dtype=scalar_type, shape=(coefficient.ufl_element().value_size(),)) + funarg = lp.GlobalArg(name, dtype=scalar_type, shape=(ufl_element.value_size(),)) expression = gem.reshape(gem.Variable(name, (None,)), - coefficient.ufl_shape) + ufl_element.value_shape()) return funarg, expression - finat_element = create_element(coefficient.ufl_element()) - + finat_element = create_element(ufl_element) shape = finat_element.index_shape size = numpy.prod(shape, dtype=int) if not interior_facet: expression = gem.reshape(gem.Variable(name, (size,)), shape) else: - varexp = gem.Variable(name, (2*size,)) + varexp = gem.Variable(name, (2 * size,)) plus = gem.view(varexp, slice(size)) - minus = gem.view(varexp, slice(size, 2*size)) + minus = gem.view(varexp, slice(size, 2 * size)) expression = (gem.reshape(plus, shape), gem.reshape(minus, shape)) size = size * 2 funarg = lp.GlobalArg(name, dtype=scalar_type, shape=(size,)) @@ -365,7 +418,6 @@ def prepare_arguments(arguments, multiindices, scalar_type, interior_facet=False expressions - GEM expressions referring to the argument tensor """ - assert isinstance(interior_facet, bool) if len(arguments) == 0: diff --git a/tsfc/kernel_interface/ufc.py b/tsfc/kernel_interface/ufc.py index 305e0988..87a3a30f 100644 --- a/tsfc/kernel_interface/ufc.py +++ b/tsfc/kernel_interface/ufc.py @@ -1,4 +1,5 @@ import numpy +import collections import functools from itertools import chain, product @@ -10,8 +11,9 @@ from finat import TensorFiniteElement import ufl +from ufl import MixedElement as ufl_MixedElement -from tsfc.kernel_interface.common import KernelBuilderBase +from tsfc.kernel_interface.common import KernelBuilderBase, KernelBuilderMixin, get_index_names from tsfc.finatinterface import create_element as _create_element @@ -19,20 +21,20 @@ create_element = functools.partial(_create_element, shape_innermost=False) -class KernelBuilder(KernelBuilderBase): +class KernelBuilder(KernelBuilderBase, KernelBuilderMixin): """Helper class for building a :class:`Kernel` object.""" - def __init__(self, integral_type, subdomain_id, domain_number, scalar_type=None, diagonal=False): + def __init__(self, integral_data, scalar_type, fem_scalar_type, diagonal=False): """Initialise a kernel builder.""" if diagonal: raise NotImplementedError("Assembly of diagonal not implemented yet, sorry") - super(KernelBuilder, self).__init__(scalar_type, integral_type.startswith("interior_facet")) - self.integral_type = integral_type + KernelBuilder.__init__(self, scalar_type, integral_data.integral_type.startswith("interior_facet")) + self.fem_scalar_type = fem_scalar_type - self.local_tensor = None self.coordinates_args = None self.coefficient_args = None self.coefficient_split = None + self.external_data_args = None if self.interior_facet: self._cell_orientations = (gem.Variable("cell_orientation_0", ()), @@ -40,6 +42,7 @@ def __init__(self, integral_type, subdomain_id, domain_number, scalar_type=None, else: self._cell_orientations = (gem.Variable("cell_orientation", ()),) + integral_type = integral_data.integral_type if integral_type == "exterior_facet": self._entity_number = {None: gem.VariableIndex(gem.Variable("facet", ()))} elif integral_type == "interior_facet": @@ -50,17 +53,30 @@ def __init__(self, integral_type, subdomain_id, domain_number, scalar_type=None, elif integral_type == "vertex": self._entity_number = {None: gem.VariableIndex(gem.Variable("vertex", ()))} - def set_arguments(self, arguments, multiindices): + self.set_coordinates(integral_data.domain) + self.set_cell_sizes(integral_data.domain) + self.set_coefficients(integral_data.coefficients) + + self.integral_data = integral_data + self.arguments = integral_data.arguments + self.local_tensor, self.return_variables, self.argument_multiindices = self.set_arguments(self.arguments) + self.mode_irs = collections.OrderedDict() + self.quadrature_indices = [] + + def set_arguments(self, arguments): """Process arguments. :arg arguments: :class:`ufl.Argument`s - :arg multiindices: GEM argument multiindices - :returns: GEM expression representing the return variable + :returns: :class:`coffee.Decl` function argument, + GEM expression representing the return variable, + GEM argument multiindices. """ - self.local_tensor, prepare, expressions = prepare_arguments( - arguments, multiindices, self.scalar_type, interior_facet=self.interior_facet) + argument_multiindices = tuple(create_element(arg.ufl_element()).get_indices() + for arg in arguments) + local_tensor, prepare, return_variables = prepare_arguments( + arguments, argument_multiindices, self.scalar_type, interior_facet=self.interior_facet) self.apply_glue(prepare) - return expressions + return local_tensor, return_variables, argument_multiindices def set_coordinates(self, domain): """Prepare the coordinate field. @@ -81,11 +97,10 @@ def set_cell_sizes(self, domain): """ pass - def set_coefficients(self, integral_data, form_data): + def set_coefficients(self, coefficients): """Prepare the coefficients of the form. - :arg integral_data: UFL integral data - :arg form_data: UFL form data + :arg coefficients: a tuple of `ufl.Coefficient`s. """ name = "w" self.coefficient_args = [ @@ -93,30 +108,60 @@ def set_coefficients(self, integral_data, form_data): pointers=[("const",), ()], qualifiers=["const"]) ] + for i, c in enumerate(coefficients): + expression = prepare_coefficient(c.ufl_element(), i, name, self.interior_facet) + self.coefficient_map[c] = expression - # enabled_coefficients is a boolean array that indicates which - # of reduced_coefficients the integral requires. - for n in range(len(integral_data.enabled_coefficients)): - if not integral_data.enabled_coefficients[n]: - continue + def set_external_data(self, elements): + """Prepare external data structures. - coefficient = form_data.function_replace_map[form_data.reduced_coefficients[n]] - expression = prepare_coefficient(coefficient, n, name, self.interior_facet) - self.coefficient_map[coefficient] = expression + :arg elements: a tuple of `ufl.FiniteElement`s. + :returns: gem expressions for the data represented by elements. - def construct_kernel(self, name, impero_c, index_names, quadrature_rule=None): + The retuned gem expressions are to be used in the operations + applied to the gem expressions obtained by compiling UFL before + compiling gem. The users are responsible for bridging these + gem expressions and actual data by setting correct values in + `external_data_numbers` and `external_data_parts` in the kernel. + """ + name = "e" + self.external_data_args = [ + coffee.Decl(self.scalar_type, coffee.Symbol(name), + pointers=[("const",), ()], + qualifiers=["const"]) + ] + if any(type(element) == ufl_MixedElement for element in elements): + raise ValueError("Unable to handle `MixedElement`s.") + expressions = [] + for i, element in enumerate(elements): + expression = prepare_coefficient(element, i, name, self.interior_facet) + expressions.append(expression) + return tuple(expressions) + + def register_requirements(self, ir): + """Inspect what is referenced by the IR that needs to be + provided by the kernel interface.""" + return None, None, None + + def construct_kernel(self, name, external_data_numbers=(), external_data_parts=(), quadrature_rule=None): """Construct a fully built kernel function. This function contains the logic for building the argument list for assembly kernels. :arg name: function name - :arg impero_c: ImperoC tuple with Impero AST and other data - :arg index_names: pre-assigned index names + :arg external_data_numbers: ignored + :arg external_data_parts: ignored :arg quadrature rule: quadrature rule (not used, stubbed out for Themis integration) :returns: a COFFEE function definition object """ from tsfc.coffee import generate as generate_coffee + + impero_c, oriented, needs_cell_sizes, tabulations = self.compile_gem() + if impero_c is None: + return self.construct_empty_kernel(name) + index_cache = self.fem_config['index_cache'] + index_names = get_index_names(self.quadrature_indices, self.argument_multiindices, index_cache) body = generate_coffee(impero_c, index_names, scalar_type=self.scalar_type) return self._construct_kernel_from_body(name, body) @@ -132,16 +177,17 @@ def _construct_kernel_from_body(self, name, body, quadrature_rule): :returns: a COFFEE function definition object """ args = [self.local_tensor] - args.extend(self.coefficient_args) + args.extend(self.coefficient_args + self.external_data_args) args.extend(self.coordinates_args) # Facet and vertex number(s) - if self.integral_type == "exterior_facet": + integral_type = self.integral_data.integral_type + if integral_type == "exterior_facet": args.append(coffee.Decl("std::size_t", coffee.Symbol("facet"))) - elif self.integral_type == "interior_facet": + elif integral_type == "interior_facet": args.append(coffee.Decl("std::size_t", coffee.Symbol("facet_0"))) args.append(coffee.Decl("std::size_t", coffee.Symbol("facet_1"))) - elif self.integral_type == "vertex": + elif integral_type == "vertex": args.append(coffee.Decl("std::size_t", coffee.Symbol("vertex"))) # Cell orientation(s) @@ -170,24 +216,25 @@ def create_element(self, element, **kwargs): return create_element(element, **kwargs) -def prepare_coefficient(coefficient, num, name, interior_facet=False): +def prepare_coefficient(ufl_element, num, name, interior_facet=False): """Bridges the kernel interface and the GEM abstraction for Coefficients. - :arg coefficient: UFL Coefficient - :arg num: coefficient index in the original form + :arg ufl_element: UFL FiniteElement + :arg num: index in the original form of the coefficient that + this ufl_element is associated with :arg name: unique name to refer to the Coefficient in the kernel :arg interior_facet: interior facet integral? :returns: GEM expression referring to the Coefficient value """ varexp = gem.Variable(name, (None, None)) - if coefficient.ufl_element().family() == 'Real': - size = numpy.prod(coefficient.ufl_shape, dtype=int) + if ufl_element.family() == 'Real': + size = numpy.prod(ufl_element.value_shape(), dtype=int) data = gem.view(varexp, slice(num, num + 1), slice(size)) - return gem.reshape(data, (), coefficient.ufl_shape) + return gem.reshape(data, (), ufl_element.value_shape()) - element = create_element(coefficient.ufl_element()) + element = create_element(ufl_element) size = numpy.prod(element.index_shape, dtype=int) def expression(data):