def pointwise_expression_kernel(exprs, scalar_type): """Compile a kernel for pointwise expressions. :arg exprs: List of expressions, all on the same iteration set. :arg scalar_type: Default scalar type (numpy.dtype). :returns: a PyOP2 kernel for evaluation of the expressions.""" if len(set(e.lvalue.node_set for e in exprs)) > 1: raise ValueError("All expressions must have same node layout.") translator = Translator() assignments = tuple(compile_to_gem(expr, translator) for expr in exprs) prefix_ordering = tuple(OrderedDict.fromkeys(itertools.chain.from_iterable( node.index_ordering() for node in gem_traversal([v for v, _ in assignments]) if isinstance(node, gem.Indexed)))) impero_c = compile_gem(assignments, prefix_ordering=prefix_ordering, remove_zeros=False, emit_return_accumulate=False) coefficients = translator.varmapping args = [] plargs = [] for expr in exprs: for c, arg in zip(expr.coefficients, expr.args): try: var = coefficients.pop(c) except KeyError: continue plargs.append(arg) is_input = arg.access in [op2.INC, op2.MAX, op2.MIN, op2.READ, op2.RW] is_output = arg.access in [op2.INC, op2.MAX, op2.MIN, op2.RW, op2.WRITE] args.append(loopy.GlobalArg(var.name, shape=var.shape, dtype=c.dat.dtype, is_input=is_input, is_output=is_output)) assert len(coefficients) == 0 name = "expression_kernel" knl = generate(impero_c, args, scalar_type, kernel_name=name, return_increments=False) return firedrake.op2.Kernel(knl, name), plargs
def transfer_kernel(Pk, P1): """Compile a kernel that will map between Pk and P1. :returns: a PyOP2 kernel. The prolongation maps a solution in P1 into Pk using the natural embedding. The restriction maps a residual in the dual of Pk into the dual of P1 (it is the dual of the prolongation), computed using linearity of the test function. """ # Mapping of a residual in Pk into a residual in P1 from coffee import base as coffee from tsfc.coffee import generate as generate_coffee, SCALAR_TYPE from tsfc.parameters import default_parameters from gem import gem, impero_utils as imp # Pk should be at least the same size as P1 assert Pk.finat_element.space_dimension() >= P1.finat_element.space_dimension() # In the general case we should compute this by doing: # numpy.linalg.solve(Pkmass, PkP1mass) Pke = Pk.finat_element._element P1e = P1.finat_element._element # TODO, rework to use finat. matrix = numpy.dot(Pke.dual.to_riesz(P1e.get_nodal_basis()), P1e.get_coeffs().T).T Vout, Vin = P1, Pk weights = gem.Literal(matrix) name = "Pk_P1_mapper" funargs = [] assert Vin.shape == Vout.shape shape = (P1e.space_dimension() * Vout.value_size, Pke.space_dimension() * Vin.value_size) outarg = coffee.Decl(SCALAR_TYPE, coffee.Symbol("A", rank=shape)) i = gem.Index() j = gem.Index() k = gem.Index() indices = i, j, k A = gem.Variable("A", shape) outgem = [gem.Indexed(gem.reshape(A, (P1e.space_dimension(), Vout.value_size), (Pke.space_dimension(), Vin.value_size)), (i, k, j, k))] funargs.append(outarg) expr = gem.Indexed(weights, (i, j)) outgem, = imp.preprocess_gem(outgem) ir = imp.compile_gem([(outgem, expr)], indices) index_names = [(i, "i"), (j, "j"), (k, "k")] body = generate_coffee(ir, index_names, default_parameters()["precision"]) function = coffee.FunDecl("void", name, funargs, body, pred=["static", "inline"]) return op2.Kernel(function, name=function.name)
def transfer_kernel(Pk, P1): """Compile a kernel that will map between Pk and P1. :returns: a PyOP2 kernel. The prolongation maps a solution in P1 into Pk using the natural embedding. The restriction maps a residual in the dual of Pk into the dual of P1 (it is the dual of the prolongation), computed using linearity of the test function. """ # Mapping of a residual in Pk into a residual in P1 from coffee import base as coffee from tsfc.coffee import generate as generate_coffee, SCALAR_TYPE from tsfc.parameters import default_parameters from gem import gem, impero_utils as imp # Pk should be at least the same size as P1 assert Pk.finat_element.space_dimension() >= P1.finat_element.space_dimension() # In the general case we should compute this by doing: # numpy.linalg.solve(Pkmass, PkP1mass) Pke = Pk.finat_element._element P1e = P1.finat_element._element # TODO, rework to use finat. matrix = numpy.dot(Pke.dual.to_riesz(P1e.get_nodal_basis()), P1e.get_coeffs().T).T Vout, Vin = P1, Pk weights = gem.Literal(matrix) name = "Pk_P1_mapper" funargs = [] assert Vin.shape == Vout.shape shape = (P1e.space_dimension() * Vout.value_size, Pke.space_dimension() * Vin.value_size) outarg = coffee.Decl(SCALAR_TYPE, coffee.Symbol("A", rank=shape)) i = gem.Index() j = gem.Index() k = gem.Index() indices = i, j, k A = gem.Variable("A", shape) outgem = [gem.Indexed(gem.reshape(A, (P1e.space_dimension(), Vout.value_size), (Pke.space_dimension(), Vin.value_size)), (i, k, j, k))] funargs.append(outarg) expr = gem.Indexed(weights, (i, j)) outgem, = imp.preprocess_gem(outgem) ir = imp.compile_gem([(outgem, expr)], indices) index_names = [(i, "i"), (j, "j"), (k, "k")] body = generate_coffee(ir, index_names, default_parameters()["precision"]) function = coffee.FunDecl("void", name, funargs, body, pred=["static", "inline"]) return op2.Kernel(function, name=function.name)
def gem_to_loopy(gem_expr, var2terminal, scalar_type): """ Method encapsulating stage 2. Converts the gem expression dag into imperoc first, and then further into loopy. :return slate_loopy: 2-tuple of loopy kernel for slate operations and loopy GlobalArg for the output variable. """ # Creation of return variables for outer loopy shape = gem_expr.shape if len(gem_expr.shape) != 0 else (1, ) idx = make_indices(len(shape)) indexed_gem_expr = gem.Indexed(gem_expr, idx) args = ([ loopy.GlobalArg("output", shape=shape, dtype=scalar_type, is_output=True, is_input=True) ] + [ loopy.GlobalArg(var.name, shape=var.shape, dtype=scalar_type) for var in var2terminal.keys() ]) ret_vars = [gem.Indexed(gem.Variable("output", shape), idx)] preprocessed_gem_expr = impero_utils.preprocess_gem([indexed_gem_expr]) # glue assignments to return variable assignments = list(zip(ret_vars, preprocessed_gem_expr)) # Part A: slate to impero_c impero_c = impero_utils.compile_gem(assignments, (), remove_zeros=False) # Part B: impero_c to loopy return generate_loopy(impero_c, args, scalar_type, "slate_loopy", []), args[0].copy()
def to_reference_coordinates(ufl_coordinate_element): # Set up UFL form cell = ufl_coordinate_element.cell() domain = ufl.Mesh(ufl_coordinate_element) K = ufl.JacobianInverse(domain) x = ufl.SpatialCoordinate(domain) x0_element = ufl.VectorElement("Real", cell, 0) x0 = ufl.Coefficient(ufl.FunctionSpace(domain, x0_element)) expr = ufl.dot(K, x - x0) # Translation to GEM C = ufl_utils.coordinate_coefficient(domain) expr = ufl_utils.preprocess_expression(expr) expr = ufl_utils.replace_coordinates(expr, C) expr = ufl_utils.simplify_abs(expr) builder = firedrake_interface.KernelBuilderBase() builder._coefficient(C, "C") builder._coefficient(x0, "x0") dim = cell.topological_dimension() point = gem.Variable('X', (dim, )) context = tsfc.fem.GemPointContext( interface=builder, ufl_cell=cell, precision=parameters["precision"], point_indices=(), point_expr=point, ) translator = tsfc.fem.Translator(context) ir = map_expr_dag(translator, expr) # Unroll result ir = [gem.Indexed(ir, alpha) for alpha in numpy.ndindex(ir.shape)] # Unroll IndexSums max_extent = parameters["unroll_indexsum"] if max_extent: def predicate(index): return index.extent <= max_extent ir = gem.optimise.unroll_indexsum(ir, predicate=predicate) # Translate to COFFEE ir = impero_utils.preprocess_gem(ir) return_variable = gem.Variable('dX', (dim, )) assignments = [(gem.Indexed(return_variable, (i, )), e) for i, e in enumerate(ir)] impero_c = impero_utils.compile_gem(assignments, ()) body = tsfc.coffee.generate(impero_c, {}, parameters["precision"]) body.open_scope = False return body
def to_reference_coordinates(ufl_coordinate_element, parameters): # Set up UFL form cell = ufl_coordinate_element.cell() domain = ufl.Mesh(ufl_coordinate_element) K = ufl.JacobianInverse(domain) x = ufl.SpatialCoordinate(domain) x0_element = ufl.VectorElement("Real", cell, 0) x0 = ufl.Coefficient(ufl.FunctionSpace(domain, x0_element)) expr = ufl.dot(K, x - x0) # Translation to GEM C = ufl.Coefficient(ufl.FunctionSpace(domain, ufl_coordinate_element)) expr = ufl_utils.preprocess_expression(expr) expr = ufl_utils.simplify_abs(expr) builder = firedrake_interface.KernelBuilderBase() builder.domain_coordinate[domain] = C builder._coefficient(C, "C") builder._coefficient(x0, "x0") dim = cell.topological_dimension() point = gem.Variable('X', (dim,)) context = tsfc.fem.GemPointContext( interface=builder, ufl_cell=cell, precision=parameters["precision"], point_indices=(), point_expr=point, ) translator = tsfc.fem.Translator(context) ir = map_expr_dag(translator, expr) # Unroll result ir = [gem.Indexed(ir, alpha) for alpha in numpy.ndindex(ir.shape)] # Unroll IndexSums max_extent = parameters["unroll_indexsum"] if max_extent: def predicate(index): return index.extent <= max_extent ir = gem.optimise.unroll_indexsum(ir, predicate=predicate) # Translate to COFFEE ir = impero_utils.preprocess_gem(ir) return_variable = gem.Variable('dX', (dim,)) assignments = [(gem.Indexed(return_variable, (i,)), e) for i, e in enumerate(ir)] impero_c = impero_utils.compile_gem(assignments, ()) body = tsfc.coffee.generate(impero_c, {}, parameters["precision"]) body.open_scope = False return body
def gem_to_loopy(gem_expr): """ Method encapsulating stage 2. Converts the gem expression dag into imperoc first, and then further into loopy. :return slate_loopy: loopy kernel for slate operations. """ # Creation of return variables for outer loopy shape = gem_expr.shape if len(gem_expr.shape) != 0 else (1,) idx = make_indices(len(shape)) indexed_gem_expr = gem.Indexed(gem_expr, idx) arg = [loopy.GlobalArg("output", shape=shape)] ret_vars = [gem.Indexed(gem.Variable("output", shape), idx)] preprocessed_gem_expr = impero_utils.preprocess_gem([indexed_gem_expr]) # glue assignments to return variable assignments = list(zip(ret_vars, preprocessed_gem_expr)) # Part A: slate to impero_c impero_c = impero_utils.compile_gem(assignments, (), remove_zeros=False) # Part B: impero_c to loopy return generate_loopy(impero_c, arg, parameters["form_compiler"]["scalar_type"], "slate_loopy", [])
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_ufl_kernel(expression, to_pts, to_element, fs): import collections from ufl.algorithms.apply_function_pullbacks import apply_function_pullbacks from ufl.algorithms.apply_algebra_lowering import apply_algebra_lowering from ufl.algorithms.apply_derivatives import apply_derivatives from ufl.algorithms.apply_geometry_lowering import apply_geometry_lowering from ufl.algorithms import extract_arguments, extract_coefficients from gem import gem, impero_utils from tsfc import fem, ufl_utils from tsfc.coffee import generate as generate_coffee from tsfc.kernel_interface import (KernelBuilderBase, needs_cell_orientations, cell_orientations_coffee_arg) # Imitate the compute_form_data processing pipeline # # Unfortunately, we cannot call compute_form_data here, since # we only have an expression, not a form expression = apply_algebra_lowering(expression) expression = apply_derivatives(expression) expression = apply_function_pullbacks(expression) expression = apply_geometry_lowering(expression) expression = apply_derivatives(expression) expression = apply_geometry_lowering(expression) expression = apply_derivatives(expression) # Replace coordinates (if any) if expression.ufl_domain(): assert fs.mesh() == expression.ufl_domain() expression = ufl_utils.replace_coordinates(expression, fs.mesh().coordinates) if extract_arguments(expression): return ValueError("Cannot interpolate UFL expression with Arguments!") builder = KernelBuilderBase() args = [] coefficients = extract_coefficients(expression) for i, coefficient in enumerate(coefficients): args.append(builder.coefficient(coefficient, "w_%d" % i)) point_index = gem.Index(name='p') ir = fem.process('cell', fs.mesh().ufl_cell(), to_pts, None, point_index, (), expression, builder.coefficient_mapper, collections.defaultdict(gem.Index)) assert len(ir) == 1 # Deal with non-scalar expressions tensor_indices = () if fs.shape: tensor_indices = tuple(gem.Index() for s in fs.shape) ir = [gem.Indexed(ir[0], tensor_indices)] # Build kernel body return_var = gem.Variable('A', (len(to_pts),) + fs.shape) return_expr = gem.Indexed(return_var, (point_index,) + tensor_indices) impero_c = impero_utils.compile_gem([return_expr], ir, [point_index]) body = generate_coffee(impero_c, index_names={point_index: 'p'}) oriented = needs_cell_orientations(ir) if oriented: args.insert(0, cell_orientations_coffee_arg) # Build kernel args.insert(0, ast.Decl("double", ast.Symbol('A', rank=(len(to_pts),) + fs.shape))) kernel_code = builder.construct_kernel("expression_kernel", args, body) return op2.Kernel(kernel_code, kernel_code.name), oriented, coefficients
def compile_ufl_kernel(expression, to_pts, to_element, fs): import collections from ufl.algorithms.apply_function_pullbacks import apply_function_pullbacks from ufl.algorithms.apply_algebra_lowering import apply_algebra_lowering from ufl.algorithms.apply_derivatives import apply_derivatives from ufl.algorithms.apply_geometry_lowering import apply_geometry_lowering from ufl.algorithms import extract_arguments, extract_coefficients from gem import gem, impero_utils from tsfc import fem, ufl_utils from tsfc.coffee import generate as generate_coffee from tsfc.kernel_interface import (KernelBuilderBase, needs_cell_orientations, cell_orientations_coffee_arg) # Imitate the compute_form_data processing pipeline # # Unfortunately, we cannot call compute_form_data here, since # we only have an expression, not a form expression = apply_algebra_lowering(expression) expression = apply_derivatives(expression) expression = apply_function_pullbacks(expression) expression = apply_geometry_lowering(expression) expression = apply_derivatives(expression) expression = apply_geometry_lowering(expression) expression = apply_derivatives(expression) # Replace coordinates (if any) if expression.ufl_domain(): assert fs.mesh() == expression.ufl_domain() expression = ufl_utils.replace_coordinates(expression, fs.mesh().coordinates) if extract_arguments(expression): return ValueError("Cannot interpolate UFL expression with Arguments!") builder = KernelBuilderBase() args = [] coefficients = extract_coefficients(expression) for i, coefficient in enumerate(coefficients): args.append(builder.coefficient(coefficient, "w_%d" % i)) point_index = gem.Index(name='p') ir = fem.compile_ufl(expression, cell=fs.mesh().ufl_cell(), points=to_pts, point_index=point_index, coefficient_mapper=builder.coefficient_mapper) assert len(ir) == 1 # Deal with non-scalar expressions tensor_indices = () if fs.shape: tensor_indices = tuple(gem.Index() for s in fs.shape) ir = [gem.Indexed(ir[0], tensor_indices)] # Build kernel body return_var = gem.Variable('A', (len(to_pts),) + fs.shape) return_expr = gem.Indexed(return_var, (point_index,) + tensor_indices) impero_c = impero_utils.compile_gem([return_expr], ir, [point_index]) body = generate_coffee(impero_c, index_names={point_index: 'p'}) oriented = needs_cell_orientations(ir) if oriented: args.insert(0, cell_orientations_coffee_arg) # Build kernel args.insert(0, ast.Decl("double", ast.Symbol('A', rank=(len(to_pts),) + fs.shape))) kernel_code = builder.construct_kernel("expression_kernel", args, body) return op2.Kernel(kernel_code, kernel_code.name), oriented, coefficients
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() 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_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() 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 dg_injection_kernel(Vf, Vc, ncell): from firedrake import Tensor, AssembledVector, TestFunction, TrialFunction from firedrake.slate.slac import compile_expression macro_builder = MacroKernelBuilder(ncell) f = ufl.Coefficient(Vf) macro_builder.set_coefficients([f]) macro_builder.set_coordinates(Vf.mesh()) Vfe = create_element(Vf.ufl_element()) macro_quadrature_rule = make_quadrature(Vfe.cell, estimate_total_polynomial_degree(ufl.inner(f, f))) index_cache = {} parameters = default_parameters() integration_dim, entity_ids = lower_integral_type(Vfe.cell, "cell") macro_cfg = dict(interface=macro_builder, ufl_cell=Vf.ufl_cell(), precision=parameters["precision"], integration_dim=integration_dim, entity_ids=entity_ids, index_cache=index_cache, quadrature_rule=macro_quadrature_rule) fexpr, = fem.compile_ufl(f, **macro_cfg) X = ufl.SpatialCoordinate(Vf.mesh()) C_a, = fem.compile_ufl(X, **macro_cfg) detJ = ufl_utils.preprocess_expression(abs(ufl.JacobianDeterminant(f.ufl_domain()))) macro_detJ, = fem.compile_ufl(detJ, **macro_cfg) Vce = create_element(Vc.ufl_element()) coarse_builder = firedrake_interface.KernelBuilder("cell", "otherwise", 0) coarse_builder.set_coordinates(Vc.mesh()) argument_multiindices = (Vce.get_indices(), ) argument_multiindex, = argument_multiindices return_variable, = coarse_builder.set_arguments((ufl.TestFunction(Vc), ), argument_multiindices) integration_dim, entity_ids = lower_integral_type(Vce.cell, "cell") # Midpoint quadrature for jacobian on coarse cell. quadrature_rule = make_quadrature(Vce.cell, 0) coarse_cfg = dict(interface=coarse_builder, ufl_cell=Vc.ufl_cell(), precision=parameters["precision"], integration_dim=integration_dim, entity_ids=entity_ids, index_cache=index_cache, quadrature_rule=quadrature_rule) X = ufl.SpatialCoordinate(Vc.mesh()) K = ufl_utils.preprocess_expression(ufl.JacobianInverse(Vc.mesh())) C_0, = fem.compile_ufl(X, **coarse_cfg) K, = fem.compile_ufl(K, **coarse_cfg) i = gem.Index() j = gem.Index() C_0 = gem.Indexed(C_0, (j, )) C_0 = gem.index_sum(C_0, quadrature_rule.point_set.indices) C_a = gem.Indexed(C_a, (j, )) X_a = gem.Sum(C_0, gem.Product(gem.Literal(-1), C_a)) K_ij = gem.Indexed(K, (i, j)) K_ij = gem.index_sum(K_ij, quadrature_rule.point_set.indices) X_a = gem.index_sum(gem.Product(K_ij, X_a), (j, )) C_0, = quadrature_rule.point_set.points C_0 = gem.Indexed(gem.Literal(C_0), (i, )) # fine quad points in coarse reference space. X_a = gem.Sum(C_0, gem.Product(gem.Literal(-1), X_a)) X_a = gem.ComponentTensor(X_a, (i, )) # Coarse basis function evaluated at fine quadrature points phi_c = fem.fiat_to_ufl(Vce.point_evaluation(0, X_a, (Vce.cell.get_dimension(), 0)), 0) tensor_indices = tuple(gem.Index(extent=d) for d in f.ufl_shape) phi_c = gem.Indexed(phi_c, argument_multiindex + tensor_indices) fexpr = gem.Indexed(fexpr, tensor_indices) quadrature_weight = macro_quadrature_rule.weight_expression expr = gem.Product(gem.IndexSum(gem.Product(phi_c, fexpr), tensor_indices), gem.Product(macro_detJ, quadrature_weight)) quadrature_indices = macro_builder.indices + macro_quadrature_rule.point_set.indices reps = spectral.Integrals([expr], quadrature_indices, argument_multiindices, parameters) assignments = spectral.flatten([(return_variable, reps)], index_cache) return_variables, expressions = zip(*assignments) expressions = impero_utils.preprocess_gem(expressions, **spectral.finalise_options) assignments = list(zip(return_variables, expressions)) impero_c = impero_utils.compile_gem(assignments, quadrature_indices + argument_multiindex, remove_zeros=True) 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) index_names.extend(zip(macro_builder.indices, ["entity"])) body = generate_coffee(impero_c, index_names, parameters["precision"]) retarg = ast.Decl(SCALAR_TYPE, ast.Symbol("R", rank=(Vce.space_dimension(), ))) local_tensor = coarse_builder.local_tensor local_tensor.init = ast.ArrayInit(numpy.zeros(Vce.space_dimension(), dtype=SCALAR_TYPE)) body.children.insert(0, local_tensor) args = [retarg] + macro_builder.kernel_args + [macro_builder.coordinates_arg, coarse_builder.coordinates_arg] # Now we have the kernel that computes <f, phi_c>dx_c # So now we need to hit it with the inverse mass matrix on dx_c u = TrialFunction(Vc) v = TestFunction(Vc) expr = Tensor(ufl.inner(u, v)*ufl.dx).inv * AssembledVector(ufl.Coefficient(Vc)) Ainv, = compile_expression(expr) Ainv = Ainv.kinfo.kernel A = ast.Symbol(local_tensor.sym.symbol) R = ast.Symbol("R") body.children.append(ast.FunCall(Ainv.name, R, coarse_builder.coordinates_arg.sym, A)) from coffee.base import Node assert isinstance(Ainv._code, Node) return op2.Kernel(ast.Node([Ainv._code, ast.FunDecl("void", "pyop2_kernel_injection_dg", args, body, pred=["static", "inline"])]), name="pyop2_kernel_injection_dg", cpp=True, include_dirs=Ainv._include_dirs, headers=Ainv._headers)
def dg_injection_kernel(Vf, Vc, ncell): from firedrake import Tensor, AssembledVector, TestFunction, TrialFunction from firedrake.slate.slac import compile_expression macro_builder = MacroKernelBuilder(ScalarType_c, ncell) f = ufl.Coefficient(Vf) macro_builder.set_coefficients([f]) macro_builder.set_coordinates(Vf.mesh()) Vfe = create_element(Vf.ufl_element()) macro_quadrature_rule = make_quadrature( Vfe.cell, estimate_total_polynomial_degree(ufl.inner(f, f))) index_cache = {} parameters = default_parameters() integration_dim, entity_ids = lower_integral_type(Vfe.cell, "cell") macro_cfg = dict(interface=macro_builder, ufl_cell=Vf.ufl_cell(), precision=parameters["precision"], integration_dim=integration_dim, entity_ids=entity_ids, index_cache=index_cache, quadrature_rule=macro_quadrature_rule) fexpr, = fem.compile_ufl(f, **macro_cfg) X = ufl.SpatialCoordinate(Vf.mesh()) C_a, = fem.compile_ufl(X, **macro_cfg) detJ = ufl_utils.preprocess_expression( abs(ufl.JacobianDeterminant(f.ufl_domain()))) macro_detJ, = fem.compile_ufl(detJ, **macro_cfg) Vce = create_element(Vc.ufl_element()) coarse_builder = firedrake_interface.KernelBuilder("cell", "otherwise", 0, ScalarType_c) coarse_builder.set_coordinates(Vc.mesh()) argument_multiindices = (Vce.get_indices(), ) argument_multiindex, = argument_multiindices return_variable, = coarse_builder.set_arguments((ufl.TestFunction(Vc), ), argument_multiindices) integration_dim, entity_ids = lower_integral_type(Vce.cell, "cell") # Midpoint quadrature for jacobian on coarse cell. quadrature_rule = make_quadrature(Vce.cell, 0) coarse_cfg = dict(interface=coarse_builder, ufl_cell=Vc.ufl_cell(), precision=parameters["precision"], integration_dim=integration_dim, entity_ids=entity_ids, index_cache=index_cache, quadrature_rule=quadrature_rule) X = ufl.SpatialCoordinate(Vc.mesh()) K = ufl_utils.preprocess_expression(ufl.JacobianInverse(Vc.mesh())) C_0, = fem.compile_ufl(X, **coarse_cfg) K, = fem.compile_ufl(K, **coarse_cfg) i = gem.Index() j = gem.Index() C_0 = gem.Indexed(C_0, (j, )) C_0 = gem.index_sum(C_0, quadrature_rule.point_set.indices) C_a = gem.Indexed(C_a, (j, )) X_a = gem.Sum(C_0, gem.Product(gem.Literal(-1), C_a)) K_ij = gem.Indexed(K, (i, j)) K_ij = gem.index_sum(K_ij, quadrature_rule.point_set.indices) X_a = gem.index_sum(gem.Product(K_ij, X_a), (j, )) C_0, = quadrature_rule.point_set.points C_0 = gem.Indexed(gem.Literal(C_0), (i, )) # fine quad points in coarse reference space. X_a = gem.Sum(C_0, gem.Product(gem.Literal(-1), X_a)) X_a = gem.ComponentTensor(X_a, (i, )) # Coarse basis function evaluated at fine quadrature points phi_c = fem.fiat_to_ufl( Vce.point_evaluation(0, X_a, (Vce.cell.get_dimension(), 0)), 0) tensor_indices = tuple(gem.Index(extent=d) for d in f.ufl_shape) phi_c = gem.Indexed(phi_c, argument_multiindex + tensor_indices) fexpr = gem.Indexed(fexpr, tensor_indices) quadrature_weight = macro_quadrature_rule.weight_expression expr = gem.Product(gem.IndexSum(gem.Product(phi_c, fexpr), tensor_indices), gem.Product(macro_detJ, quadrature_weight)) quadrature_indices = macro_builder.indices + macro_quadrature_rule.point_set.indices reps = spectral.Integrals([expr], quadrature_indices, argument_multiindices, parameters) assignments = spectral.flatten([(return_variable, reps)], index_cache) return_variables, expressions = zip(*assignments) expressions = impero_utils.preprocess_gem(expressions, **spectral.finalise_options) assignments = list(zip(return_variables, expressions)) impero_c = impero_utils.compile_gem(assignments, quadrature_indices + argument_multiindex, remove_zeros=True) 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) index_names.extend(zip(macro_builder.indices, ["entity"])) body = generate_coffee(impero_c, index_names, parameters["precision"], ScalarType_c) retarg = ast.Decl(ScalarType_c, ast.Symbol("R", rank=(Vce.space_dimension(), ))) local_tensor = coarse_builder.local_tensor local_tensor.init = ast.ArrayInit( numpy.zeros(Vce.space_dimension(), dtype=ScalarType_c)) body.children.insert(0, local_tensor) args = [retarg] + macro_builder.kernel_args + [ macro_builder.coordinates_arg, coarse_builder.coordinates_arg ] # Now we have the kernel that computes <f, phi_c>dx_c # So now we need to hit it with the inverse mass matrix on dx_c u = TrialFunction(Vc) v = TestFunction(Vc) expr = Tensor(ufl.inner(u, v) * ufl.dx).inv * AssembledVector( ufl.Coefficient(Vc)) Ainv, = compile_expression(expr) Ainv = Ainv.kinfo.kernel A = ast.Symbol(local_tensor.sym.symbol) R = ast.Symbol("R") body.children.append( ast.FunCall(Ainv.name, R, coarse_builder.coordinates_arg.sym, A)) from coffee.base import Node assert isinstance(Ainv._code, Node) return op2.Kernel(ast.Node([ Ainv._code, ast.FunDecl("void", "pyop2_kernel_injection_dg", args, body, pred=["static", "inline"]) ]), name="pyop2_kernel_injection_dg", cpp=True, include_dirs=Ainv._include_dirs, headers=Ainv._headers)
def gencode(expr): impero_c = impero_utils.compile_gem([Ri], [expr], (i, j)) return impero_c.tree
def gencode(expr): impero_c = impero_utils.compile_gem([(Ri, expr)], (i, j)) return impero_c.tree
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 gencode(expr): impero_c = impero_utils.compile_gem([Ri], [expr], (i, j)) return impero_c.tree
def gencode(expr): impero_c = impero_utils.compile_gem([(Ri, expr)], (i, j)) return impero_c.tree
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)