def compile_expression_at_points(expression, points, coordinates, parameters=None): """Compiles a UFL expression to be evaluated at compile-time known reference points. Useful for interpolating UFL expressions onto function spaces with only point evaluation nodes. :arg expression: UFL expression :arg points: reference coordinates of the evaluation points :arg coordinates: the coordinate function :arg parameters: parameters object """ import coffee.base as ast if parameters is None: parameters = default_parameters() else: _ = default_parameters() _.update(parameters) parameters = _ # No arguments, please! if extract_arguments(expression): return ValueError("Cannot interpolate UFL expression with Arguments!") # Apply UFL preprocessing expression = ufl_utils.preprocess_expression(expression) # Initialise kernel builder builder = firedrake_interface.ExpressionKernelBuilder() # Replace coordinates (if any) domain = expression.ufl_domain() if domain: assert coordinates.ufl_domain() == domain builder.domain_coordinate[domain] = coordinates # Collect required coefficients coefficients = extract_coefficients(expression) if has_type(expression, GeometricQuantity): coefficients = [coordinates] + coefficients builder.set_coefficients(coefficients) # Split mixed coefficients expression = ufl_utils.split_coefficients(expression, builder.coefficient_split) # Translate to GEM point_set = PointSet(points) config = dict(interface=builder, ufl_cell=coordinates.ufl_domain().ufl_cell(), precision=parameters["precision"], point_set=point_set) ir, = fem.compile_ufl(expression, point_sum=False, **config) # Deal with non-scalar expressions value_shape = ir.shape tensor_indices = tuple(gem.Index() for s in value_shape) if value_shape: ir = gem.Indexed(ir, tensor_indices) # Build kernel body return_shape = (len(points),) + value_shape return_indices = point_set.indices + tensor_indices return_var = gem.Variable('A', return_shape) return_arg = ast.Decl(SCALAR_TYPE, ast.Symbol('A', rank=return_shape)) return_expr = gem.Indexed(return_var, return_indices) ir, = impero_utils.preprocess_gem([ir]) impero_c = impero_utils.compile_gem([(return_expr, ir)], return_indices) point_index, = point_set.indices body = generate_coffee(impero_c, {point_index: 'p'}, parameters["precision"]) # Handle cell orientations if builder.needs_cell_orientations([ir]): builder.require_cell_orientations() # Build kernel tuple return builder.construct_kernel(return_arg, body)
def compile_integral(integral_data, form_data, prefix, parameters, interface=firedrake_interface): """Compiles a UFL integral into an assembly kernel. :arg integral_data: UFL integral data :arg form_data: UFL form data :arg prefix: kernel name will start with this string :arg parameters: parameters object :arg interface: backend module for the kernel interface :returns: a kernel constructed by the kernel interface """ if parameters is None: parameters = default_parameters() else: _ = default_parameters() _.update(parameters) parameters = _ # 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.KernelBuilder(integral_type, integral_data.subdomain_id, domain_numbering[integral_data.domain]) argument_multiindices = tuple(builder.create_element(arg.ufl_element()).get_indices() for arg in arguments) return_variables = builder.set_arguments(arguments, argument_multiindices) builder.set_coordinates(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, precision=parameters["precision"], integration_dim=integration_dim, entity_ids=entity_ids, argument_multiindices=argument_multiindices, index_cache=index_cache) mode_irs = collections.OrderedDict() 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, 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 for COFFEE 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)) # Look for cell orientations in the IR if builder.needs_cell_orientations(expressions): builder.require_cell_orientations() # 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) # Construct kernel body = generate_coffee(impero_c, index_names, parameters["precision"], expressions, split_argument_indices) return builder.construct_kernel(kernel_name, body)
def compile_integral(integral_data, form_data, prefix, parameters, interface=firedrake_interface): """Compiles a UFL integral into an assembly kernel. :arg integral_data: UFL integral data :arg form_data: UFL form data :arg prefix: kernel name will start with this string :arg parameters: parameters object :arg interface: backend module for the kernel interface :returns: a kernel constructed by the kernel interface """ if parameters is None: parameters = default_parameters() else: _ = default_parameters() _.update(parameters) parameters = _ # 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() fiat_cell = as_fiat_cell(cell) integration_dim, entity_ids = lower_integral_type(fiat_cell, integral_type) argument_indices = tuple(tuple(gem.Index(extent=e) for e in create_element(arg.ufl_element()).index_shape) for arg in arguments) flat_argument_indices = tuple(chain(*argument_indices)) quadrature_indices = [] # Dict mapping domains to index in original_form.ufl_domains() domain_numbering = form_data.original_form.domain_numbering() builder = interface.KernelBuilder(integral_type, integral_data.subdomain_id, domain_numbering[integral_data.domain]) return_variables = builder.set_arguments(arguments, argument_indices) coordinates = ufl_utils.coordinate_coefficient(mesh) builder.set_coordinates(coordinates) builder.set_coefficients(integral_data, form_data) # Map from UFL FiniteElement objects to Index instances. This is # so we reuse Index instances when evaluating the same coefficient # multiple times with the same table. Occurs, for example, if we # have multiple integrals here (and the affine coordinate # evaluation can be hoisted). index_cache = collections.defaultdict(gem.Index) kernel_cfg = dict(interface=builder, ufl_cell=cell, precision=parameters["precision"], integration_dim=integration_dim, entity_ids=entity_ids, argument_indices=argument_indices, index_cache=index_cache) kernel_cfg["facetarea"] = facetarea_generator(mesh, coordinates, kernel_cfg, integral_type) kernel_cfg["cellvolume"] = cellvolume_generator(mesh, coordinates, kernel_cfg) irs = [] for integral in integral_data.integrals: params = {} # Record per-integral parameters params.update(integral.metadata()) if params.get("quadrature_rule") == "default": del params["quadrature_rule"] # parameters override per-integral metadata params.update(parameters) integrand = ufl_utils.replace_coordinates(integral.integrand(), coordinates) 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: 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 += quadrature_multiindex config = kernel_cfg.copy() config.update(quadrature_rule=quad_rule) ir = fem.compile_ufl(integrand, interior_facet=interior_facet, **config) if parameters["unroll_indexsum"]: def predicate(index): return index.extent <= parameters["unroll_indexsum"] ir = opt.unroll_indexsum(ir, predicate=predicate) ir = [gem.index_sum(expr, quadrature_multiindex) for expr in ir] irs.append(ir) # Sum the expressions that are part of the same restriction ir = list(reduce(gem.Sum, e, gem.Zero()) for e in zip(*irs)) # Need optimised roots for COFFEE ir = impero_utils.preprocess_gem(ir) # Look for cell orientations in the IR if builder.needs_cell_orientations(ir): builder.require_cell_orientations() impero_c = impero_utils.compile_gem(return_variables, ir, tuple(quadrature_indices) + flat_argument_indices, remove_zeros=True) # Generate COFFEE index_names = [(si, name + str(n)) for index, name in zip(argument_indices, ['j', 'k']) for n, si in enumerate(index)] if len(quadrature_indices) == 1: index_names.append((quadrature_indices[0], 'ip')) else: for i, quadrature_index in enumerate(quadrature_indices): index_names.append((quadrature_index, 'ip_%d' % i)) body = generate_coffee(impero_c, index_names, parameters["precision"], ir, flat_argument_indices) kernel_name = "%s_%s_integral_%s" % (prefix, integral_type, integral_data.subdomain_id) return builder.construct_kernel(kernel_name, body)
def compile_integral(integral_data, form_data, prefix, parameters, interface=firedrake_interface): """Compiles a UFL integral into an assembly kernel. :arg integral_data: UFL integral data :arg form_data: UFL form data :arg prefix: kernel name will start with this string :arg parameters: parameters object :arg interface: backend module for the kernel interface :returns: a kernel constructed by the kernel interface """ if parameters is None: parameters = default_parameters() else: _ = default_parameters() _.update(parameters) parameters = _ # 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) fiat_cell = as_fiat_cell(cell) integration_dim, entity_ids = lower_integral_type(fiat_cell, integral_type) argument_multiindices = tuple( create_element(arg.ufl_element()).get_indices() for arg in arguments) argument_indices = tuple(chain(*argument_multiindices)) quadrature_indices = [] # Dict mapping domains to index in original_form.ufl_domains() domain_numbering = form_data.original_form.domain_numbering() builder = interface.KernelBuilder(integral_type, integral_data.subdomain_id, domain_numbering[integral_data.domain]) return_variables = builder.set_arguments(arguments, argument_multiindices) coordinates = ufl_utils.coordinate_coefficient(mesh) builder.set_coordinates(coordinates) 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. index_cache = {} kernel_cfg = dict(interface=builder, ufl_cell=cell, precision=parameters["precision"], integration_dim=integration_dim, entity_ids=entity_ids, argument_multiindices=argument_multiindices, index_cache=index_cache) kernel_cfg["facetarea"] = facetarea_generator(mesh, coordinates, kernel_cfg, integral_type) kernel_cfg["cellvolume"] = cellvolume_generator(mesh, coordinates, kernel_cfg) mode_irs = collections.OrderedDict() 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_utils.replace_coordinates(integral.integrand(), coordinates) integrand = ufl.replace(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: 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, 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 iteritems(mode_irs): assignments.extend(mode.flatten(viewitems(var_reps))) if assignments: return_variables, expressions = zip(*assignments) else: return_variables = [] expressions = [] # Need optimised roots for COFFEE options = dict( reduce( operator.and_, [viewitems(mode.finalise_options) for mode in iterkeys(mode_irs)])) expressions = impero_utils.preprocess_gem(expressions, **options) assignments = list(zip(return_variables, expressions)) # Look for cell orientations in the IR if builder.needs_cell_orientations(expressions): builder.require_cell_orientations() # Construct ImperoC index_ordering = tuple(quadrature_indices) + 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 = [(si, name + str(n)) for index, name in zip(argument_multiindices, ['j', 'k']) for n, si in enumerate(index)] if len(quadrature_indices) == 1: index_names.append((quadrature_indices[0], 'ip')) else: for i, quadrature_index in enumerate(quadrature_indices): index_names.append((quadrature_index, 'ip_%d' % i)) # Construct kernel body = generate_coffee(impero_c, index_names, parameters["precision"], expressions, argument_indices) return builder.construct_kernel(kernel_name, body)
def compile_expression_at_points(expression, points, coordinates, interface=None, parameters=None, coffee=True): """Compiles a UFL expression to be evaluated at compile-time known reference points. Useful for interpolating UFL expressions onto function spaces with only point evaluation nodes. :arg expression: UFL expression :arg points: reference coordinates of the evaluation points :arg coordinates: the coordinate function :arg interface: backend module for the kernel interface :arg parameters: parameters object :arg coffee: compile coffee kernel instead of loopy kernel """ import coffee.base as ast import loopy as lp if parameters is None: parameters = default_parameters() else: _ = default_parameters() _.update(parameters) parameters = _ # Determine whether in complex mode complex_mode = is_complex(parameters["scalar_type"]) # Apply UFL preprocessing expression = ufl_utils.preprocess_expression(expression, complex_mode=complex_mode) # Initialise kernel builder if interface is None: if coffee: import tsfc.kernel_interface.firedrake as firedrake_interface_coffee interface = firedrake_interface_coffee.ExpressionKernelBuilder else: # Delayed import, loopy is a runtime dependency import tsfc.kernel_interface.firedrake_loopy as firedrake_interface_loopy interface = firedrake_interface_loopy.ExpressionKernelBuilder builder = interface(parameters["scalar_type"]) arguments = extract_arguments(expression) argument_multiindices = tuple( builder.create_element(arg.ufl_element()).get_indices() for arg in arguments) # Replace coordinates (if any) domain = expression.ufl_domain() if domain: assert coordinates.ufl_domain() == domain builder.domain_coordinate[domain] = coordinates builder.set_cell_sizes(domain) # Collect required coefficients coefficients = extract_coefficients(expression) if has_type(expression, GeometricQuantity) or any( fem.needs_coordinate_mapping(c.ufl_element()) for c in coefficients): coefficients = [coordinates] + coefficients builder.set_coefficients(coefficients) # Split mixed coefficients expression = ufl_utils.split_coefficients(expression, builder.coefficient_split) # Translate to GEM point_set = PointSet(points) config = dict(interface=builder, ufl_cell=coordinates.ufl_domain().ufl_cell(), precision=parameters["precision"], point_set=point_set, argument_multiindices=argument_multiindices) ir, = fem.compile_ufl(expression, point_sum=False, **config) # Deal with non-scalar expressions value_shape = ir.shape tensor_indices = tuple(gem.Index() for s in value_shape) if value_shape: ir = gem.Indexed(ir, tensor_indices) # Build kernel body return_indices = point_set.indices + tensor_indices + tuple( chain(*argument_multiindices)) return_shape = tuple(i.extent for i in return_indices) return_var = gem.Variable('A', return_shape) if coffee: return_arg = ast.Decl(parameters["scalar_type"], ast.Symbol('A', rank=return_shape)) else: return_arg = lp.GlobalArg("A", dtype=parameters["scalar_type"], shape=return_shape) return_expr = gem.Indexed(return_var, return_indices) ir, = impero_utils.preprocess_gem([ir]) impero_c = impero_utils.compile_gem([(return_expr, ir)], return_indices) point_index, = point_set.indices # Handle kernel interface requirements builder.register_requirements([ir]) # Build kernel tuple return builder.construct_kernel(return_arg, impero_c, parameters["precision"], {point_index: 'p'})
def compile_expression_dual_evaluation(expression, to_element, *, domain=None, interface=None, parameters=None, coffee=False): """Compile a UFL expression to be evaluated against a compile-time known reference element's dual basis. Useful for interpolating UFL expressions into e.g. N1curl spaces. :arg expression: UFL expression :arg to_element: A FInAT element for the target space :arg domain: optional UFL domain the expression is defined on (required when expression contains no domain). :arg interface: backend module for the kernel interface :arg parameters: parameters object :arg coffee: compile coffee kernel instead of loopy kernel """ import coffee.base as ast import loopy as lp # Just convert FInAT element to FIAT for now. # Dual evaluation in FInAT will bring a thorough revision. to_element = to_element.fiat_equivalent if any(len(dual.deriv_dict) != 0 for dual in to_element.dual_basis()): raise NotImplementedError( "Can only interpolate onto dual basis functionals without derivative evaluation, sorry!" ) if parameters is None: parameters = default_parameters() else: _ = default_parameters() _.update(parameters) parameters = _ # Determine whether in complex mode complex_mode = is_complex(parameters["scalar_type"]) # Find out which mapping to apply try: mapping, = set(to_element.mapping()) except ValueError: raise NotImplementedError( "Don't know how to interpolate onto zany spaces, sorry") expression = apply_mapping(expression, mapping, domain) # Apply UFL preprocessing expression = ufl_utils.preprocess_expression(expression, complex_mode=complex_mode) # Initialise kernel builder if interface is None: if coffee: import tsfc.kernel_interface.firedrake as firedrake_interface_coffee interface = firedrake_interface_coffee.ExpressionKernelBuilder else: # Delayed import, loopy is a runtime dependency import tsfc.kernel_interface.firedrake_loopy as firedrake_interface_loopy interface = firedrake_interface_loopy.ExpressionKernelBuilder builder = interface(parameters["scalar_type"]) arguments = extract_arguments(expression) argument_multiindices = tuple( builder.create_element(arg.ufl_element()).get_indices() for arg in arguments) # Replace coordinates (if any) unless otherwise specified by kwarg if domain is None: domain = expression.ufl_domain() assert domain is not None # Collect required coefficients first_coefficient_fake_coords = False coefficients = extract_coefficients(expression) if has_type(expression, GeometricQuantity) or any( fem.needs_coordinate_mapping(c.ufl_element()) for c in coefficients): # Create a fake coordinate coefficient for a domain. coords_coefficient = ufl.Coefficient( ufl.FunctionSpace(domain, domain.ufl_coordinate_element())) builder.domain_coordinate[domain] = coords_coefficient builder.set_cell_sizes(domain) coefficients = [coords_coefficient] + coefficients first_coefficient_fake_coords = True builder.set_coefficients(coefficients) # Split mixed coefficients expression = ufl_utils.split_coefficients(expression, builder.coefficient_split) # Translate to GEM kernel_cfg = dict( interface=builder, ufl_cell=domain.ufl_cell(), # FIXME: change if we ever implement # interpolation on facets. integral_type="cell", argument_multiindices=argument_multiindices, index_cache={}, scalar_type=parameters["scalar_type"]) if all( isinstance(dual, PointEvaluation) for dual in to_element.dual_basis()): # This is an optimisation for point-evaluation nodes which # should go away once FInAT offers the interface properly qpoints = [] # Everything is just a point evaluation. for dual in to_element.dual_basis(): ptdict = dual.get_point_dict() qpoint, = ptdict.keys() (qweight, component), = ptdict[qpoint] assert allclose(qweight, 1.0) assert component == () qpoints.append(qpoint) point_set = PointSet(qpoints) config = kernel_cfg.copy() config.update(point_set=point_set) # Allow interpolation onto QuadratureElements to refer to the quadrature # rule they represent if isinstance(to_element, FIAT.QuadratureElement): assert allclose(asarray(qpoints), asarray(to_element._points)) quad_rule = QuadratureRule(point_set, to_element._weights) config["quadrature_rule"] = quad_rule expr, = fem.compile_ufl(expression, **config, point_sum=False) # In some cases point_set.indices may be dropped from expr, but nothing # new should now appear assert set(expr.free_indices) <= set( chain(point_set.indices, *argument_multiindices)) shape_indices = tuple(gem.Index() for _ in expr.shape) basis_indices = point_set.indices ir = gem.Indexed(expr, shape_indices) else: # This is general code but is more unrolled than necssary. dual_expressions = [] # one for each functional broadcast_shape = len(expression.ufl_shape) - len( to_element.value_shape()) shape_indices = tuple(gem.Index() for _ in expression.ufl_shape[:broadcast_shape]) expr_cache = {} # Sharing of evaluation of the expression at points for dual in to_element.dual_basis(): pts = tuple(sorted(dual.get_point_dict().keys())) try: expr, point_set = expr_cache[pts] except KeyError: point_set = PointSet(pts) config = kernel_cfg.copy() config.update(point_set=point_set) expr, = fem.compile_ufl(expression, **config, point_sum=False) # In some cases point_set.indices may be dropped from expr, but # nothing new should now appear assert set(expr.free_indices) <= set( chain(point_set.indices, *argument_multiindices)) expr = gem.partial_indexed(expr, shape_indices) expr_cache[pts] = expr, point_set weights = collections.defaultdict(list) for p in pts: for (w, cmp) in dual.get_point_dict()[p]: weights[cmp].append(w) qexprs = gem.Zero() for cmp in sorted(weights): qweights = gem.Literal(weights[cmp]) qexpr = gem.Indexed(expr, cmp) qexpr = gem.index_sum( gem.Indexed(qweights, point_set.indices) * qexpr, point_set.indices) qexprs = gem.Sum(qexprs, qexpr) assert qexprs.shape == () assert set(qexprs.free_indices) == set( chain(shape_indices, *argument_multiindices)) dual_expressions.append(qexprs) basis_indices = (gem.Index(), ) ir = gem.Indexed(gem.ListTensor(dual_expressions), basis_indices) # Build kernel body return_indices = basis_indices + shape_indices + tuple( chain(*argument_multiindices)) return_shape = tuple(i.extent for i in return_indices) return_var = gem.Variable('A', return_shape) if coffee: return_arg = ast.Decl(parameters["scalar_type"], ast.Symbol('A', rank=return_shape)) else: return_arg = lp.GlobalArg("A", dtype=parameters["scalar_type"], shape=return_shape) return_expr = gem.Indexed(return_var, return_indices) # TODO: one should apply some GEM optimisations as in assembly, # but we don't for now. ir, = impero_utils.preprocess_gem([ir]) impero_c = impero_utils.compile_gem([(return_expr, ir)], return_indices) index_names = dict( (idx, "p%d" % i) for (i, idx) in enumerate(basis_indices)) # Handle kernel interface requirements builder.register_requirements([ir]) # Build kernel tuple return builder.construct_kernel(return_arg, impero_c, index_names, first_coefficient_fake_coords)
def compile_integral(integral_data, form_data, prefix, parameters, interface=firedrake_interface): """Compiles a UFL integral into an assembly kernel. :arg integral_data: UFL integral data :arg form_data: UFL form data :arg prefix: kernel name will start with this string :arg parameters: parameters object :arg interface: backend module for the kernel interface :returns: a kernel constructed by the kernel interface """ if parameters is None: parameters = default_parameters() else: _ = default_parameters() _.update(parameters) parameters = _ # 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() fiat_cell = as_fiat_cell(cell) integration_dim, entity_ids = lower_integral_type(fiat_cell, integral_type) argument_indices = tuple(gem.Index(name=name) for arg, name in zip(arguments, ['j', 'k'])) quadrature_indices = [] # Dict mapping domains to index in original_form.ufl_domains() domain_numbering = form_data.original_form.domain_numbering() builder = interface.KernelBuilder(integral_type, integral_data.subdomain_id, domain_numbering[integral_data.domain]) return_variables = builder.set_arguments(arguments, argument_indices) coordinates = ufl_utils.coordinate_coefficient(mesh) if ufl_utils.is_element_affine(mesh.ufl_coordinate_element()): # For affine mesh geometries we prefer code generation that # composes well with optimisations. builder.set_coordinates(coordinates, mode='list_tensor') else: # Otherwise we use the approach that might be faster (?) builder.set_coordinates(coordinates) builder.set_coefficients(integral_data, form_data) # Map from UFL FiniteElement objects to Index instances. This is # so we reuse Index instances when evaluating the same coefficient # multiple times with the same table. Occurs, for example, if we # have multiple integrals here (and the affine coordinate # evaluation can be hoisted). index_cache = collections.defaultdict(gem.Index) kernel_cfg = dict(interface=builder, ufl_cell=cell, precision=parameters["precision"], integration_dim=integration_dim, entity_ids=entity_ids, argument_indices=argument_indices, index_cache=index_cache) kernel_cfg["facetarea"] = facetarea_generator(mesh, coordinates, kernel_cfg, integral_type) kernel_cfg["cellvolume"] = cellvolume_generator(mesh, coordinates, kernel_cfg) irs = [] for integral in integral_data.integrals: params = {} # Record per-integral parameters params.update(integral.metadata()) if params.get("quadrature_rule") == "default": del params["quadrature_rule"] # parameters override per-integral metadata params.update(parameters) # 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"]) integration_cell = fiat_cell.construct_subelement(integration_dim) quad_rule = params.get("quadrature_rule", create_quadrature(integration_cell, quadrature_degree)) if not isinstance(quad_rule, QuadratureRule): raise ValueError("Expected to find a QuadratureRule object, not a %s" % type(quad_rule)) integrand = ufl_utils.replace_coordinates(integral.integrand(), coordinates) integrand = ufl_utils.split_coefficients(integrand, builder.coefficient_split) quadrature_index = gem.Index(name='ip') quadrature_indices.append(quadrature_index) config = kernel_cfg.copy() config.update(quadrature_rule=quad_rule, point_index=quadrature_index) ir = fem.compile_ufl(integrand, interior_facet=interior_facet, **config) if parameters["unroll_indexsum"]: ir = opt.unroll_indexsum(ir, max_extent=parameters["unroll_indexsum"]) irs.append([(gem.IndexSum(expr, quadrature_index) if quadrature_index in expr.free_indices else expr) for expr in ir]) # Sum the expressions that are part of the same restriction ir = list(reduce(gem.Sum, e, gem.Zero()) for e in zip(*irs)) # Need optimised roots for COFFEE ir = opt.remove_componenttensors(ir) # Look for cell orientations in the IR if builder.needs_cell_orientations(ir): builder.require_cell_orientations() impero_c = impero_utils.compile_gem(return_variables, ir, tuple(quadrature_indices) + argument_indices, remove_zeros=True) # Generate COFFEE index_names = [(index, index.name) for index in argument_indices] if len(quadrature_indices) == 1: index_names.append((quadrature_indices[0], 'ip')) else: for i, quadrature_index in enumerate(quadrature_indices): index_names.append((quadrature_index, 'ip_%d' % i)) body = generate_coffee(impero_c, index_names, parameters["precision"], ir, argument_indices) kernel_name = "%s_%s_integral_%s" % (prefix, integral_type, integral_data.subdomain_id) return builder.construct_kernel(kernel_name, body)