diff --git a/tsfc/driver.py b/tsfc/driver.py index 21db6d9a..d69bf848 100644 --- a/tsfc/driver.py +++ b/tsfc/driver.py @@ -1,31 +1,21 @@ import collections -import operator -import string import time import sys -from functools import reduce from itertools import chain from finat.physically_mapped import DirectlyDefinedElement, PhysicallyMappedElement -from numpy import asarray - 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.log import GREEN -from ufl.utils.sequences import max_degree import gem import gem.impero_utils as impero_utils -from FIAT.reference_element import TensorProductCell - import finat -from finat.quadrature import AbstractQuadratureRule, make_quadrature 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 @@ -34,6 +24,28 @@ sys.setrecursionlimit(3000) +TSFCIntegralDataInfo = collections.namedtuple("TSFCIntegralDataInfo", + ["domain", "integral_type", "subdomain_id", "domain_number", + "arguments", + "coefficients", "coefficient_numbers"]) +TSFCIntegralDataInfo.__doc__ = """ + Minimal set of objects for kernel builders. + + domain - The mesh. + integral_type - The type of integral. + subdomain_id - What is the subdomain id for this kernel. + domain_number - Which domain number in the original form + does this kernel correspond to (can be used to index into + original_form.ufl_domains() to get the correct domain). + coefficients - A list of coefficients. + coefficient_numbers - A list of which coefficients from the + form the kernel needs. + + This is a minimal set of objects that kernel builders need to + construct a kernel from :attr:`integrals` of :class:`~ufl.IntegralData`. + """ + + def compile_form(form, prefix="form", parameters=None, interface=None, coffee=True, diagonal=False): """Compiles a UFL form into a set of assembly kernels. @@ -76,12 +88,7 @@ def compile_integral(integral_data, form_data, prefix, parameters, interface, co :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 = _ + parameters = preprocess_parameters(parameters) if interface is None: if coffee: import tsfc.kernel_interface.firedrake as firedrake_interface_coffee @@ -90,180 +97,61 @@ 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"] - + scalar_type = parameters["scalar_type"] 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], + domain_number = domain_numbering[integral_data.domain] + coefficients = [form_data.function_replace_map[c] for c in integral_data.integral_coefficients] + # 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 = tuple(form_data.original_coefficient_positions[i] + for i, (_, enabled) in enumerate(zip(form_data.reduced_coefficients, integral_data.enabled_coefficients)) + if enabled) + integral_data_info = TSFCIntegralDataInfo(domain=integral_data.domain, + integral_type=integral_data.integral_type, + subdomain_id=integral_data.subdomain_id, + domain_number=domain_number, + arguments=arguments, + coefficients=coefficients, + coefficient_numbers=coefficient_numbers) + builder = interface(integral_data_info, 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() + ctx = builder.create_context() for integral in 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) - - config = kernel_cfg.copy() - config.update(quadrature_rule=quad_rule) - expressions = fem.compile_ufl(integrand, - fem.PointSetContext(**config), - interior_facet=interior_facet) - 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)) + integrand_exprs = builder.compile_integrand(integrand, params, ctx) + integral_exprs = builder.construct_integrals(integrand_exprs, params) + builder.stash_integrals(integral_exprs, params, ctx) + return builder.construct_kernel(kernel_name, ctx) - 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, ufl_element, *, @@ -429,71 +317,3 @@ def __call__(self, ps): assert set(gem_expr.free_indices) <= set(chain(ps.indices, *argument_multiindices)) return gem_expr - - -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..fa88e678 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 + +from ufl.utils.sequences import max_degree import coffee.base as coffee +from FIAT.reference_element import TensorProductCell + +from finat.quadrature import AbstractQuadratureRule, make_quadrature + 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 class KernelBuilderBase(KernelInterface): @@ -107,3 +124,302 @@ def register_requirements(self, ir): """ # Nothing is required by default pass + + +class KernelBuilderMixin(object): + """Mixin for KernelBuilder classes.""" + + def compile_integrand(self, integrand, params, ctx): + """Compile UFL integrand. + + :arg integrand: UFL integrand. + :arg params: a dict containing "quadrature_rule". + :arg ctx: context created with :meth:`create_context` method. + + See :meth:`create_context` for typical calling sequence. + """ + # Split Coefficients + if self.coefficient_split: + integrand = ufl_utils.split_coefficients(integrand, self.coefficient_split) + # Compile: ufl -> gem + info = self.integral_data_info + functions = list(info.arguments) + [self.coordinate(info.domain)] + list(info.coefficients) + set_quad_rule(params, info.domain.ufl_cell(), info.integral_type, functions) + quad_rule = params["quadrature_rule"] + config = self.fem_config() + config['argument_multiindices'] = self.argument_multiindices + config['quadrature_rule'] = quad_rule + config['index_cache'] = ctx['index_cache'] + expressions = fem.compile_ufl(integrand, + fem.PointSetContext(**config), + interior_facet=self.interior_facet) + ctx['quadrature_indices'].extend(quad_rule.point_set.indices) + return expressions + + def construct_integrals(self, integrand_expressions, params): + """Construct integrals from integrand expressions. + + :arg integrand_expressions: gem expressions for integrands. + :arg params: a dict containing "mode" and "quadrature_rule". + + integrand_expressions must be indexed with :attr:`argument_multiindices`; + these gem expressions are obtained by calling :meth:`compile_integrand` + method or by modifying the gem expressions returned by + :meth:`compile_integrand`. + + See :meth:`create_context` for typical calling sequence. + """ + mode = pick_mode(params["mode"]) + return mode.Integrals(integrand_expressions, + params["quadrature_rule"].point_set.indices, + self.argument_multiindices, + params) + + def stash_integrals(self, reps, params, ctx): + """Stash integral representations in ctx. + + :arg reps: integral representations. + :arg params: a dict containing "mode". + :arg ctx: context in which reps are stored. + + See :meth:`create_context` for typical calling sequence. + """ + mode = pick_mode(params["mode"]) + mode_irs = ctx['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, ctx): + """Compile gem representation of integrals to impero_c. + + :arg ctx: the context containing the gem representation of integrals. + :returns: a tuple of impero_c, oriented, needs_cell_sizes, tabulations + required to finally construct a kernel in :meth:`construct_kernel`. + + See :meth:`create_context` for typical calling sequence. + """ + # Finalise mode representations into a set of assignments + mode_irs = ctx['mode_irs'] + + assignments = [] + for mode, var_reps in mode_irs.items(): + assignments.extend(mode.flatten(var_reps.items(), ctx['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(ctx['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 + + def fem_config(self): + """Return a dictionary used with fem.compile_ufl. + + One needs to update this dictionary with "argument_multiindices", + "quadrature_rule", and "index_cache" before using this with + fem.compile_ufl. + """ + info = self.integral_data_info + integral_type = info.integral_type + cell = info.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, + scalar_type=self.fem_scalar_type) + + def create_context(self): + """Create builder context. + + *index_cache* + + 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). + + *quadrature_indices* + + List of quadrature indices used. + + *mode_irs* + + Dict for mode representations. + + For each set of integrals to make a kernel for (i,e., + `integral_data.integrals`), one must first create a ctx object by + calling :meth:`create_context` method. + This ctx object collects objects associated with the integrals that + are eventually used to construct the kernel. + The following is a typical calling sequence: + + .. code-block:: python3 + + builder = KernelBuilder(...) + params = {"mode": "spectral"} + ctx = builder.create_context() + for integral in integral_data.integrals: + integrand = integral.integrand() + integrand_exprs = builder.compile_integrand(integrand, params, ctx) + integral_exprs = builder.construct_integrals(integrand_exprs, params) + builder.stash_integrals(integral_exprs, params, ctx) + kernel = builder.construct_kernel(kernel_name, ctx) + + """ + return {'index_cache': {}, + 'quadrature_indices': [], + 'mode_irs': collections.OrderedDict()} + + +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 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 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 25977d8e..6f17f1bb 100644 --- a/tsfc/kernel_interface/firedrake.py +++ b/tsfc/kernel_interface/firedrake.py @@ -14,7 +14,7 @@ from tsfc import kernel_args from tsfc.coffee import generate as generate_coffee 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 def make_builder(*args, **kwargs): @@ -23,7 +23,7 @@ def make_builder(*args, **kwargs): class Kernel(object): __slots__ = ("ast", "arguments", "integral_type", "oriented", "subdomain_id", - "domain_number", "needs_cell_sizes", "tabulations", "quadrature_rule", + "domain_number", "needs_cell_sizes", "tabulations", "coefficient_numbers", "name", "__weakref__", "flop_count") """A compiled Kernel object. @@ -37,16 +37,17 @@ 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 tabulations: The runtime tabulations this kernel requires :kwarg needs_cell_sizes: Does the kernel require cell sizes. :kwarg flop_count: Estimated total flops for this kernel. """ def __init__(self, ast=None, arguments=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, - flop_count=0): + tabulations=None, + flop_count=0, + name=None): # Defaults self.ast = ast self.arguments = arguments @@ -56,7 +57,9 @@ def __init__(self, ast=None, arguments=None, integral_type=None, oriented=False, self.subdomain_id = subdomain_id self.coefficient_numbers = coefficient_numbers self.needs_cell_sizes = needs_cell_sizes + self.tabulations = tabulations self.flop_count = flop_count + self.name = name super(Kernel, self).__init__() @@ -123,16 +126,16 @@ def create_element(self, element, **kwargs): return create_element(element, **kwargs) -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, + def __init__(self, integral_data_info, scalar_type, dont_split=(), diagonal=False): """Initialise a kernel builder.""" - super(KernelBuilder, self).__init__(scalar_type, integral_type.startswith("interior_facet")) + integral_type = integral_data_info.integral_type + super(KernelBuilder, self).__init__(coffee.as_cstr(scalar_type), integral_type.startswith("interior_facet")) + self.fem_scalar_type = 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 @@ -153,19 +156,31 @@ 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_arguments(integral_data_info.arguments) + self.integral_data_info = integral_data_info + + 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 """ - funarg, exprs = prepare_arguments(arguments, multiindices, - self.scalar_type, - interior_facet=self.interior_facet, - diagonal=self.diagonal) + 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) + funarg, return_variables = prepare_arguments(arguments, + argument_multiindices, + self.scalar_type, + interior_facet=self.interior_facet, + diagonal=self.diagonal) self.output_arg = kernel_args.OutputKernelArg(funarg) - return exprs + self.return_variables = return_variables + self.argument_multiindices = argument_multiindices def set_coordinates(self, domain): """Prepare the coordinate field. @@ -185,7 +200,6 @@ def set_coefficients(self, integral_data, form_data): :arg form_data: UFL form data """ 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)): @@ -203,68 +217,68 @@ def set_coefficients(self, integral_data, form_data): self.coefficient_split[coefficient] = split 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): coeff_coffee_arg = self._coefficient(coefficient, f"w_{i}") self.coefficient_args.append(kernel_args.CoefficientKernelArg(coeff_coffee_arg)) - self.kernel.coefficient_numbers = tuple(coefficient_numbers) 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, name, ctx): """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 name: kernel name + :arg ctx: kernel builder context to get impero_c from :returns: :class:`Kernel` object """ + impero_c, oriented, needs_cell_sizes, tabulations = self.compile_gem(ctx) + if impero_c is None: + return self.construct_empty_kernel(name) + info = self.integral_data_info args = [self.output_arg, self.coordinates_arg] - if self.kernel.oriented: + if oriented: ori_coffee_arg = coffee.Decl("int", coffee.Symbol("cell_orientations"), pointers=[("restrict",)], qualifiers=["const"]) args.append(kernel_args.CellOrientationsKernelArg(ori_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"]: + if info.integral_type in ["exterior_facet", "exterior_facet_vert"]: ext_coffee_arg = coffee.Decl("unsigned int", coffee.Symbol("facet", rank=(1,)), qualifiers=["const"]) args.append(kernel_args.ExteriorFacetKernelArg(ext_coffee_arg)) - elif self.kernel.integral_type in ["interior_facet", "interior_facet_vert"]: + elif info.integral_type in ["interior_facet", "interior_facet_vert"]: int_coffee_arg = coffee.Decl("unsigned int", coffee.Symbol("facet", rank=(2,)), qualifiers=["const"]) args.append(kernel_args.InteriorFacetKernelArg(int_coffee_arg)) - for n, shape in self.kernel.tabulations: + for name_, shape in tabulations: tab_coffee_arg = coffee.Decl(self.scalar_type, - coffee.Symbol(n, rank=shape), + coffee.Symbol(name_, rank=shape), qualifiers=["const"]) args.append(kernel_args.TabulationKernelArg(tab_coffee_arg)) - - coffee_args = [a.coffee_arg for a in args] + index_names = get_index_names(ctx['quadrature_indices'], self.argument_multiindices, ctx['index_cache']) body = generate_coffee(impero_c, index_names, self.scalar_type) - - self.kernel.ast = KernelBuilderBase.construct_kernel(self, name, coffee_args, body) - self.kernel.arguments = tuple(args) - self.kernel.quadrature_rule = quadrature_rule - self.kernel.name = name - self.kernel.flop_count = count_flops(impero_c) - return self.kernel + ast = KernelBuilderBase.construct_kernel(self, name, [a.coffee_arg for a in args], body) + flop_count = count_flops(impero_c) # Estimated total flops for this kernel. + return Kernel(ast=ast, + arguments=tuple(args), + integral_type=info.integral_type, + subdomain_id=info.subdomain_id, + domain_number=info.domain_number, + coefficient_numbers=info.coefficient_numbers, + oriented=oriented, + needs_cell_sizes=needs_cell_sizes, + tabulations=tabulations, + flop_count=flop_count, + name=name) def construct_empty_kernel(self, name): """Return None, since Firedrake needs no empty kernels. diff --git a/tsfc/kernel_interface/firedrake_loopy.py b/tsfc/kernel_interface/firedrake_loopy.py index 38585948..1bb35a5b 100644 --- a/tsfc/kernel_interface/firedrake_loopy.py +++ b/tsfc/kernel_interface/firedrake_loopy.py @@ -13,7 +13,7 @@ from tsfc import kernel_args 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 @@ -29,7 +29,7 @@ def make_builder(*args, **kwargs): class Kernel: __slots__ = ("ast", "arguments", "integral_type", "oriented", "subdomain_id", - "domain_number", "needs_cell_sizes", "tabulations", "quadrature_rule", + "domain_number", "needs_cell_sizes", "tabulations", "coefficient_numbers", "name", "flop_count", "__weakref__") """A compiled Kernel object. @@ -43,17 +43,18 @@ class Kernel: 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 tabulations: The runtime tabulations this kernel requires :kwarg needs_cell_sizes: Does the kernel require cell sizes. :kwarg name: The name of this kernel. :kwarg flop_count: Estimated total flops for this kernel. """ def __init__(self, ast=None, arguments=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, - flop_count=0): + tabulations=None, + flop_count=0, + name=None): # Defaults self.ast = ast self.arguments = arguments @@ -63,7 +64,9 @@ def __init__(self, ast=None, arguments=None, integral_type=None, oriented=False, self.subdomain_id = subdomain_id self.coefficient_numbers = coefficient_numbers self.needs_cell_sizes = needs_cell_sizes + self.tabulations = tabulations self.flop_count = flop_count + self.name = name class KernelBuilderBase(_KernelBuilderBase): @@ -198,16 +201,16 @@ def construct_kernel(self, impero_c, index_names, first_coefficient_fake_coords) self.tabulations, name, args, count_flops(impero_c)) -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_info, scalar_type, + dont_split=(), diagonal=False): """Initialise a kernel builder.""" - super().__init__(scalar_type, integral_type.startswith("interior_facet")) + integral_type = integral_data_info.integral_type + super(KernelBuilder, self).__init__(scalar_type, integral_type.startswith("interior_facet")) + self.fem_scalar_type = 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 @@ -228,18 +231,31 @@ 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_arguments(integral_data_info.arguments) + self.integral_data_info = integral_data_info + + 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 """ - funarg, expressions = prepare_arguments( - arguments, multiindices, self.scalar_type, interior_facet=self.interior_facet, - diagonal=self.diagonal) + 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) + funarg, return_variables = prepare_arguments(arguments, + argument_multiindices, + self.scalar_type, + interior_facet=self.interior_facet, + diagonal=self.diagonal) self.output_arg = kernel_args.OutputKernelArg(funarg) - return expressions + self.return_variables = return_variables + self.argument_multiindices = argument_multiindices def set_coordinates(self, domain): """Prepare the coordinate field. @@ -259,7 +275,6 @@ def set_coefficients(self, integral_data, form_data): :arg form_data: UFL form data """ 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)): @@ -277,57 +292,59 @@ def set_coefficients(self, integral_data, form_data): self.coefficient_split[coefficient] = split 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): coeff_loopy_arg = self._coefficient(coefficient, f"w_{i}") self.coefficient_args.append(kernel_args.CoefficientKernelArg(coeff_loopy_arg)) - self.kernel.coefficient_numbers = tuple(coefficient_numbers) 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, name, ctx): """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 name: kernel name + :arg ctx: kernel builder context to get impero_c from :returns: :class:`Kernel` object """ + impero_c, oriented, needs_cell_sizes, tabulations = self.compile_gem(ctx) + if impero_c is None: + return self.construct_empty_kernel(name) + info = self.integral_data_info args = [self.output_arg, self.coordinates_arg] - if self.kernel.oriented: + if oriented: args.append(self.cell_orientations_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"]: + if info.integral_type in ["exterior_facet", "exterior_facet_vert"]: ext_loopy_arg = lp.GlobalArg("facet", numpy.uint32, shape=(1,)) args.append(kernel_args.ExteriorFacetKernelArg(ext_loopy_arg)) - elif self.kernel.integral_type in ["interior_facet", "interior_facet_vert"]: + elif info.integral_type in ["interior_facet", "interior_facet_vert"]: int_loopy_arg = lp.GlobalArg("facet", numpy.uint32, shape=(2,)) args.append(kernel_args.InteriorFacetKernelArg(int_loopy_arg)) - for name_, shape in self.kernel.tabulations: + for name_, shape in tabulations: tab_loopy_arg = lp.GlobalArg(name_, dtype=self.scalar_type, shape=shape) args.append(kernel_args.TabulationKernelArg(tab_loopy_arg)) - - self.kernel.ast = generate_loopy(impero_c, [arg.loopy_arg for arg in args], - self.scalar_type, name, index_names) - self.kernel.arguments = tuple(args) - self.kernel.name = name - self.kernel.quadrature_rule = quadrature_rule - self.kernel.flop_count = count_flops(impero_c) - return self.kernel + index_names = get_index_names(ctx['quadrature_indices'], self.argument_multiindices, ctx['index_cache']) + ast = generate_loopy(impero_c, [arg.loopy_arg for arg in args], + self.scalar_type, name, index_names) + flop_count = count_flops(impero_c) # Estimated total flops for this kernel. + return Kernel(ast=ast, + arguments=tuple(args), + integral_type=info.integral_type, + subdomain_id=info.subdomain_id, + domain_number=info.domain_number, + coefficient_numbers=info.coefficient_numbers, + oriented=oriented, + needs_cell_sizes=needs_cell_sizes, + tabulations=tabulations, + flop_count=flop_count, + name=name) def construct_empty_kernel(self, name): """Return None, since Firedrake needs no empty kernels. diff --git a/tsfc/kernel_interface/ufc.py b/tsfc/kernel_interface/ufc.py index 305e0988..65c91e51 100644 --- a/tsfc/kernel_interface/ufc.py +++ b/tsfc/kernel_interface/ufc.py @@ -11,7 +11,7 @@ import ufl -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,14 +19,16 @@ 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_info, scalar_type, diagonal=False): """Initialise a kernel builder.""" + integral_type = integral_data_info.integral_type if diagonal: raise NotImplementedError("Assembly of diagonal not implemented yet, sorry") super(KernelBuilder, self).__init__(scalar_type, integral_type.startswith("interior_facet")) + self.fem_scalar_type = scalar_type self.integral_type = integral_type self.local_tensor = None @@ -50,17 +52,32 @@ 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_arguments(integral_data_info.arguments) + self.integral_data_info = integral_data_info + + 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 """ - 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) + 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, prepare, return_variables = prepare_arguments(arguments, + argument_multiindices, + self.scalar_type, + interior_facet=self.interior_facet) self.apply_glue(prepare) - return expressions + self.local_tensor = local_tensor + self.return_variables = return_variables + self.argument_multiindices = argument_multiindices def set_coordinates(self, domain): """Prepare the coordinate field. @@ -104,23 +121,32 @@ def set_coefficients(self, integral_data, form_data): expression = prepare_coefficient(coefficient, n, name, self.interior_facet) self.coefficient_map[coefficient] = expression - def construct_kernel(self, name, impero_c, index_names, quadrature_rule=None): + 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, ctx): """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 name: kernel name + :arg ctx: kernel builder context to get impero_c from :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, _, _, _ = self.compile_gem(ctx) + if impero_c is None: + return self.construct_empty_kernel(name) + index_names = get_index_names(ctx['quadrature_indices'], self.argument_multiindices, ctx['index_cache']) body = generate_coffee(impero_c, index_names, scalar_type=self.scalar_type) return self._construct_kernel_from_body(name, body) - def _construct_kernel_from_body(self, name, body, quadrature_rule): + def _construct_kernel_from_body(self, name, body): """Construct a fully built kernel function. This function contains the logic for building the argument