def extruded_int_horiz_facet(exp, builder, top_sks, bottom_sks, coordsym, mesh_layer_sym, cell_orientations): """Generates a code statement for evaluating interior horizontal facet integrals. :arg exp: A :class:`TensorBase` expression. :arg builder: A :class:`KernelBuilder` containing the expression context. :arg top_sks: An iterable of index ordered TSFC kernels for the top kernels. :arg bottom_sks: An iterable of index ordered TSFC kernels for the bottom kernels. :arg coordsym: An `ast.Symbol` object representing coordinate arguments for the kernel. :arg mesh_layer_sym: An `ast.Symbol` representing the mesh layer. :arg cell_orientations: An `ast.Symbol` representing cell orientation information. Returns: A COFFEE code statement and updated include_dirs """ t = builder.temps[exp] nlayers = exp.ufl_domain().topological.layers - 1 incl = [] top_calls = [] bottom_calls = [] for top, btm in zip(top_sks, bottom_sks): assert top.indices == btm.indices, ( "Top and bottom kernels must have the same indices" ) index = top.indices # Generate an iterable of coefficients to pass to the subkernel # if any are required c_set = top.kinfo.coefficient_map + btm.kinfo.coefficient_map coefficient_map = tuple(OrderedDict.fromkeys(c_set)) clist = [c for ci in coefficient_map for c in builder.coefficient(exp.coefficients()[ci])] # TODO: Is this safe? if top.kinfo.oriented and btm.kinfo.oriented: clist.append(cell_orientations) dirs = top.kinfo.kernel._include_dirs + btm.kinfo.kernel._include_dirs incl.extend(tuple(OrderedDict.fromkeys(dirs))) tensor = eigen_tensor(exp, t, index) top_calls.append(ast.FunCall(top.kinfo.kernel.name, tensor, coordsym, *clist)) bottom_calls.append(ast.FunCall(btm.kinfo.kernel.name, tensor, coordsym, *clist)) else_stmt = ast.Block(top_calls + bottom_calls, open_scope=True) inter_stmt = ast.If(ast.Eq(mesh_layer_sym, nlayers - 1), (ast.Block(bottom_calls, open_scope=True), else_stmt)) stmt = ast.If(ast.Eq(mesh_layer_sym, 0), (ast.Block(top_calls, open_scope=True), inter_stmt)) return stmt, incl
def extruded_top_bottom_facet(cxt_kernel, builder, coordsym, mesh_layer_sym, cell_orientations): """Generates a code statement for evaluating exterior top/bottom facet integrals. :arg cxt_kernel: A :namedtuple:`ContextKernel` containing all relevant integral types and TSFC kernels associated with the form nested in the expression. :arg builder: A :class:`KernelBuilder` containing the expression context. :arg coordsym: An `ast.Symbol` object representing coordinate arguments for the kernel. :arg mesh_layer_sym: An `ast.Symbol` representing the mesh layer. :arg cell_orientations: An `ast.Symbol` representing cell orientation information. Returns: A COFFEE code statement and updated include_dirs """ exp = cxt_kernel.tensor t = builder.temps[exp] nlayers = exp.ufl_domain().topological.layers - 1 incl = [] body = [] for splitkernel in cxt_kernel.tsfc_kernels: index = splitkernel.indices kinfo = splitkernel.kinfo # Generate an iterable of coefficients to pass to the subkernel # if any are required clist = [ c for ci in kinfo.coefficient_map for c in builder.coefficient(exp.coefficients()[ci]) ] if kinfo.oriented: clist.insert(0, cell_orientations) incl.extend(kinfo.kernel._include_dirs) tensor = eigen_tensor(exp, t, index) body.append(ast.FunCall(kinfo.kernel.name, tensor, coordsym, *clist)) if cxt_kernel.original_integral_type == "exterior_facet_bottom": layer = 0 else: layer = nlayers - 1 stmt = ast.If(ast.Eq(mesh_layer_sym, layer), [ast.Block(body, open_scope=True)]) return stmt, incl
def tensor_assembly_calls(builder): """Generates a block of statements for assembling the local finite element tensors. :arg builder: The :class:`LocalKernelBuilder` containing all relevant expression information and assembly calls. """ assembly_calls = builder.assembly_calls statements = [ast.FlatBlock("/* Assemble local tensors */\n")] # Cell integrals are straightforward. Just splat them out. statements.extend(assembly_calls["cell"]) if builder.needs_cell_facets: # The for-loop will have the general structure: # # FOR (facet=0; facet<num_facets; facet++): # IF (facet is interior): # *interior calls # ELSE IF (facet is exterior): # *exterior calls # # If only interior (exterior) facets are present, # then only a single IF-statement checking for interior # (exterior) facets will be present within the loop. The # cell facets are labelled `1` for interior, and `0` for # exterior. statements.append(ast.FlatBlock("/* Loop over cell facets */\n")) int_calls = list( chain(*[ assembly_calls[it_type] for it_type in ("interior_facet", "interior_facet_vert") ])) ext_calls = list( chain(*[ assembly_calls[it_type] for it_type in ("exterior_facet", "exterior_facet_vert") ])) # Generate logical statements for handling exterior/interior facet # integrals on subdomains. # Currently only facet integrals are supported. for sd_type in ("subdomains_exterior_facet", "subdomains_interior_facet"): stmts = [] for sd, sd_calls in groupby(assembly_calls[sd_type], lambda x: x[0]): _, calls = zip(*sd_calls) if_sd = ast.Eq( ast.Symbol(builder.cell_facet_sym, rank=(builder.it_sym, 1)), sd) stmts.append( ast.If(if_sd, (ast.Block(calls, open_scope=True), ))) if sd_type == "subdomains_exterior_facet": ext_calls.extend(stmts) if sd_type == "subdomains_interior_facet": int_calls.extend(stmts) # Compute the number of facets to loop over domain = builder.expression.ufl_domain() if domain.cell_set._extruded: num_facets = domain.ufl_cell()._cells[0].num_facets() else: num_facets = domain.ufl_cell().num_facets() if_ext = ast.Eq( ast.Symbol(builder.cell_facet_sym, rank=(builder.it_sym, 0)), 0) if_int = ast.Eq( ast.Symbol(builder.cell_facet_sym, rank=(builder.it_sym, 0)), 1) body = [] if ext_calls: body.append( ast.If(if_ext, (ast.Block(ext_calls, open_scope=True), ))) if int_calls: body.append( ast.If(if_int, (ast.Block(int_calls, open_scope=True), ))) statements.append( ast.For(ast.Decl("unsigned int", builder.it_sym, init=0), ast.Less(builder.it_sym, num_facets), ast.Incr(builder.it_sym, 1), body)) if builder.needs_mesh_layers: # In the presence of interior horizontal facet calls, an # IF-ELIF-ELSE block is generated using the mesh levels # as conditions for which calls are needed: # # IF (layer == bottom_layer): # *bottom calls # ELSE IF (layer == top_layer): # *top calls # ELSE: # *top calls # *bottom calls # # Any extruded top or bottom calls for extruded facets are # included within the appropriate mesh-level IF-blocks. If # no interior horizontal facet calls are present, then # standard IF-blocks are generated for exterior top/bottom # facet calls when appropriate: # # IF (layer == bottom_layer): # *bottom calls # # IF (layer == top_layer): # *top calls # # The mesh level is an integer provided as a macro kernel # argument. # FIXME: No variable layers assumption statements.append(ast.FlatBlock("/* Mesh levels: */\n")) num_layers = ast.Symbol(builder.mesh_layer_count_sym, rank=(0, )) layer = builder.mesh_layer_sym types = [ "interior_facet_horiz_top", "interior_facet_horiz_bottom", "exterior_facet_top", "exterior_facet_bottom" ] decide = [ ast.Less(layer, num_layers), ast.Greater(layer, 0), ast.Eq(layer, num_layers), ast.Eq(layer, 0) ] for (integral_type, which) in zip(types, decide): statements.append( ast.If(which, (ast.Block(assembly_calls[integral_type], open_scope=True), ))) return statements
def tensor_assembly_calls(builder): """Generates a block of statements for assembling the local finite element tensors. :arg builder: The :class:`LocalKernelBuilder` containing all relevant expression information and assembly calls. """ statements = [ast.FlatBlock("/* Assemble local tensors */\n")] # Cell integrals are straightforward. Just splat them out. statements.extend(builder.assembly_calls["cell"]) if builder.needs_cell_facets: # The for-loop will have the general structure: # # FOR (facet=0; facet<num_facets; facet++): # IF (facet is interior): # *interior calls # ELSE IF (facet is exterior): # *exterior calls # # If only interior (exterior) facets are present, # then only a single IF-statement checking for interior # (exterior) facets will be present within the loop. The # cell facets are labelled `1` for interior, and `0` for # exterior. statements.append(ast.FlatBlock("/* Loop over cell facets */\n")) int_calls = list( chain(*[ builder.assembly_calls[it_type] for it_type in ("interior_facet", "interior_facet_vert") ])) ext_calls = list( chain(*[ builder.assembly_calls[it_type] for it_type in ("exterior_facet", "exterior_facet_vert") ])) # Compute the number of facets to loop over domain = builder.expression.ufl_domain() if domain.cell_set._extruded: num_facets = domain.ufl_cell()._cells[0].num_facets() else: num_facets = domain.ufl_cell().num_facets() if_ext = ast.Eq( ast.Symbol(builder.cell_facet_sym, rank=(builder.it_sym, )), 0) if_int = ast.Eq( ast.Symbol(builder.cell_facet_sym, rank=(builder.it_sym, )), 1) body = [] if ext_calls: body.append( ast.If(if_ext, (ast.Block(ext_calls, open_scope=True), ))) if int_calls: body.append( ast.If(if_int, (ast.Block(int_calls, open_scope=True), ))) statements.append( ast.For(ast.Decl("unsigned int", builder.it_sym, init=0), ast.Less(builder.it_sym, num_facets), ast.Incr(builder.it_sym, 1), body)) if builder.needs_mesh_layers: # In the presence of interior horizontal facet calls, an # IF-ELIF-ELSE block is generated using the mesh levels # as conditions for which calls are needed: # # IF (layer == bottom_layer): # *bottom calls # ELSE IF (layer == top_layer): # *top calls # ELSE: # *top calls # *bottom calls # # Any extruded top or bottom calls for extruded facets are # included within the appropriate mesh-level IF-blocks. If # no interior horizontal facet calls are present, then # standard IF-blocks are generated for exterior top/bottom # facet calls when appropriate: # # IF (layer == bottom_layer): # *bottom calls # # IF (layer == top_layer): # *top calls # # The mesh level is an integer provided as a macro kernel # argument. # FIXME: No variable layers assumption statements.append(ast.FlatBlock("/* Mesh levels: */\n")) num_layers = builder.expression.ufl_domain().topological.layers - 1 int_top = builder.assembly_calls["interior_facet_horiz_top"] int_btm = builder.assembly_calls["interior_facet_horiz_bottom"] ext_top = builder.assembly_calls["exterior_facet_top"] ext_btm = builder.assembly_calls["exterior_facet_bottom"] bottom = ast.Block(int_top + ext_btm, open_scope=True) top = ast.Block(int_btm + ext_top, open_scope=True) rest = ast.Block(int_btm + int_top, open_scope=True) statements.append( ast.If(ast.Eq(builder.mesh_layer_sym, 0), (bottom, ast.If(ast.Eq(builder.mesh_layer_sym, num_layers - 1), (top, rest))))) return statements
def compile_expression(slate_expr, tsfc_parameters=None): """Takes a SLATE expression `slate_expr` and returns the appropriate :class:`firedrake.op2.Kernel` object representing the SLATE expression. :arg slate_expr: a :class:'TensorBase' expression. :arg tsfc_parameters: an optional `dict` of form compiler parameters to be passed onto TSFC during the compilation of ufl forms. """ if not isinstance(slate_expr, TensorBase): raise ValueError( "Expecting a `slate.TensorBase` expression, not a %r" % slate_expr) # TODO: Get PyOP2 to write into mixed dats if any(len(a.function_space()) > 1 for a in slate_expr.arguments()): raise NotImplementedError("Compiling mixed slate expressions") # Initialize shape and statements list shape = slate_expr.shape statements = [] # Create a builder for the SLATE expression builder = KernelBuilder(expression=slate_expr, tsfc_parameters=tsfc_parameters) # Initialize coordinate and facet symbols coordsym = ast.Symbol("coords") coords = None cellfacetsym = ast.Symbol("cell_facets") inc = [] # Now we construct the list of statements to provide to the builder context_temps = builder.temps.copy() for exp, t in context_temps.items(): statements.append(ast.Decl(eigen_matrixbase_type(exp.shape), t)) statements.append(ast.FlatBlock("%s.setZero();\n" % t)) for splitkernel in builder.kernel_exprs[exp]: clist = [] index = splitkernel.indices kinfo = splitkernel.kinfo integral_type = kinfo.integral_type if integral_type not in [ "cell", "interior_facet", "exterior_facet" ]: raise NotImplementedError( "Integral type %s not currently supported." % integral_type) coordinates = exp.ufl_domain().coordinates if coords is not None: assert coordinates == coords else: coords = coordinates for cindex in kinfo.coefficient_map: c = exp.coefficients()[cindex] # Handles both mixed and non-mixed coefficient cases clist.extend(builder.extract_coefficient(c)) inc.extend(kinfo.kernel._include_dirs) tensor = eigen_tensor(exp, t, index) if integral_type in ["interior_facet", "exterior_facet"]: builder.require_cell_facets() itsym = ast.Symbol("i0") clist.append(ast.FlatBlock("&%s" % itsym)) loop_body = [] nfacet = exp.ufl_domain().ufl_cell().num_facets() if integral_type == "exterior_facet": checker = 1 else: checker = 0 loop_body.append( ast.If( ast.Eq(ast.Symbol(cellfacetsym, rank=(itsym, )), checker), [ ast.Block([ ast.FunCall(kinfo.kernel.name, tensor, coordsym, *clist) ], open_scope=True) ])) loop = ast.For(ast.Decl("unsigned int", itsym, init=0), ast.Less(itsym, nfacet), ast.Incr(itsym, 1), loop_body) statements.append(loop) else: statements.append( ast.FunCall(kinfo.kernel.name, tensor, coordsym, *clist)) # Now we handle any terms that require auxiliary data (if any) if bool(builder.aux_exprs): aux_temps, aux_statements = auxiliary_information(builder) context_temps.update(aux_temps) statements.extend(aux_statements) result_sym = ast.Symbol("T%d" % len(builder.temps)) result_data_sym = ast.Symbol("A%d" % len(builder.temps)) result_type = "Eigen::Map<%s >" % eigen_matrixbase_type(shape) result = ast.Decl(SCALAR_TYPE, ast.Symbol(result_data_sym, shape)) result_statement = ast.FlatBlock( "%s %s((%s *)%s);\n" % (result_type, result_sym, SCALAR_TYPE, result_data_sym)) statements.append(result_statement) cpp_string = ast.FlatBlock( metaphrase_slate_to_cpp(slate_expr, context_temps)) statements.append(ast.Assign(result_sym, cpp_string)) # Generate arguments for the macro kernel args = [result, ast.Decl("%s **" % SCALAR_TYPE, coordsym)] for c in slate_expr.coefficients(): if isinstance(c, Constant): ctype = "%s *" % SCALAR_TYPE else: ctype = "%s **" % SCALAR_TYPE args.extend([ ast.Decl(ctype, sym_c) for sym_c in builder.extract_coefficient(c) ]) if builder.needs_cell_facets: args.append(ast.Decl("char *", cellfacetsym)) macro_kernel_name = "compile_slate" kernel_ast, oriented = builder.construct_ast( name=macro_kernel_name, args=args, statements=ast.Block(statements)) inc.extend(["%s/include/eigen3/" % d for d in PETSC_DIR]) op2kernel = op2.Kernel( kernel_ast, macro_kernel_name, cpp=True, include_dirs=inc, headers=['#include <Eigen/Dense>', '#define restrict __restrict']) assert len(slate_expr.ufl_domains()) == 1 kinfo = KernelInfo(kernel=op2kernel, integral_type="cell", oriented=oriented, subdomain_id="otherwise", domain_number=0, coefficient_map=range(len(slate_expr.coefficients())), needs_cell_facets=builder.needs_cell_facets) idx = tuple([0] * slate_expr.rank) return (SplitKernel(idx, kinfo), )
_to_sum = lambda o: ast.Sum(_ast(o[0]), _to_sum(o[1:])) if len( o) > 1 else _ast(o[0]) _to_prod = lambda o: ast.Prod(_ast(o[0]), _to_sum(o[1:])) if len( o) > 1 else _ast(o[0]) _to_aug_assign = lambda op, o: op(_ast(o[0]), _ast(o[1])) _ast_map = { MathFunction: (lambda e: ast.FunCall(e._name, *[_ast(o) for o in e.ufl_operands])), ufl.algebra.Sum: (lambda e: ast.Par(_to_sum(e.ufl_operands))), ufl.algebra.Product: (lambda e: ast.Par(_to_prod(e.ufl_operands))), ufl.algebra.Division: (lambda e: ast.Par(ast.Div(*[_ast(o) for o in e.ufl_operands]))), ufl.algebra.Abs: (lambda e: ast.FunCall("abs", _ast(e.ufl_operands[0]))), Assign: (lambda e: _to_aug_assign(e._ast, e.ufl_operands)), AugmentedAssignment: (lambda e: _to_aug_assign(e._ast, e.ufl_operands)), ufl.constantvalue.ScalarValue: (lambda e: ast.Symbol(e._value)), ufl.constantvalue.Zero: (lambda e: ast.Symbol(0)), ufl.classes.Conditional: (lambda e: ast.Ternary(*[_ast(o) for o in e.ufl_operands])), ufl.classes.EQ: (lambda e: ast.Eq(*[_ast(o) for o in e.ufl_operands])), ufl.classes.NE: (lambda e: ast.NEq(*[_ast(o) for o in e.ufl_operands])), ufl.classes.LT: (lambda e: ast.Less(*[_ast(o) for o in e.ufl_operands])), ufl.classes.LE: (lambda e: ast.LessEq(*[_ast(o) for o in e.ufl_operands])), ufl.classes.GT: (lambda e: ast.Greater(*[_ast(o) for o in e.ufl_operands])), ufl.classes.GE: (lambda e: ast.GreaterEq(*[_ast(o) for o in e.ufl_operands])) }
def facet_integral_loop(cxt_kernel, builder, coordsym, cellfacetsym, cell_orientations): """Generates a code statement for evaluating exterior/interior facet integrals. :arg cxt_kernel: A :namedtuple:`ContextKernel` containing all relevant integral types and TSFC kernels associated with the form nested in the expression. :arg builder: A :class:`KernelBuilder` containing the expression context. :arg coordsym: An `ast.Symbol` object representing coordinate arguments for the kernel. :arg cellfacetsym: An `ast.Symbol` representing the cell facets. :arg cell_orientations: An `ast.Symbol` representing cell orientation information. Returns: A COFFEE code statement and updated include_dirs """ exp = cxt_kernel.tensor t = builder.temps[exp] it_type = cxt_kernel.original_integral_type itsym = ast.Symbol("i0") chker = { "interior_facet": 1, "interior_facet_vert": 1, "exterior_facet": 0, "exterior_facet_vert": 0 } # Compute the correct number of facets for a particular facet measure if it_type in ["interior_facet", "exterior_facet"]: # Non-extruded case nfacet = exp.ufl_domain().ufl_cell().num_facets() elif it_type in ["interior_facet_vert", "exterior_facet_vert"]: # Extrusion case base_cell = exp.ufl_domain().ufl_cell()._cells[0] nfacet = base_cell.num_facets() else: raise ValueError("Integral type %s not supported." % it_type) incl = [] funcalls = [] checker = chker[it_type] for splitkernel in cxt_kernel.tsfc_kernels: index = splitkernel.indices kinfo = splitkernel.kinfo # Generate an iterable of coefficients to pass to the subkernel # if any are required clist = [ c for ci in kinfo.coefficient_map for c in builder.coefficient(exp.coefficients()[ci]) ] incl.extend(kinfo.kernel._include_dirs) tensor = eigen_tensor(exp, t, index) if kinfo.oriented: clist.insert(0, cell_orientations) clist.append(ast.FlatBlock("&%s" % itsym)) funcalls.append( ast.FunCall(kinfo.kernel.name, tensor, coordsym, *clist)) loop_body = ast.If( ast.Eq(ast.Symbol(cellfacetsym, rank=(itsym, )), checker), [ast.Block(funcalls, open_scope=True)]) loop_stmt = ast.For(ast.Decl("unsigned int", itsym, init=0), ast.Less(itsym, nfacet), ast.Incr(itsym, 1), loop_body) return loop_stmt, incl