def auxiliary_information(builder): """This function generates any auxiliary information regarding special handling of expressions that do not have any integral forms or subkernels associated with it. :arg builder: a :class:`SlateKernelBuilder` object that contains all the necessary temporary and expression information. Returns: a mapping of the form ``{aux_node: aux_temp}``, where `aux_node` is an already assembled data-object provided as a `ufl.Coefficient` and `aux_temp` is the corresponding temporary. a list of auxiliary statements are returned that contain temporary declarations and any code-blocks needed to evaluate the expression. """ aux_temps = {} aux_statements = [] for i, exp in enumerate(builder.aux_exprs): if isinstance(exp, Action): acting_coefficient = exp._acting_coefficient assert isinstance(acting_coefficient, Coefficient) temp = ast.Symbol("C%d" % i) V = acting_coefficient.function_space() node_extent = V.fiat_element.space_dimension() dof_extent = np.prod(V.ufl_element().value_shape()) aux_statements.append( ast.Decl( eigen_matrixbase_type(shape=(dof_extent * node_extent, )), temp)) aux_statements.append(ast.FlatBlock("%s.setZero();\n" % temp)) # Now we unpack the coefficient and insert its entries into a 1D vector temporary isym = ast.Symbol("i1") jsym = ast.Symbol("j1") tensor_index = ast.Sum(ast.Prod(dof_extent, isym), jsym) # Inner-loop running over dof_extent inner_loop = ast.For( ast.Decl("unsigned int", jsym, init=0), ast.Less(jsym, dof_extent), ast.Incr(jsym, 1), ast.Assign( ast.Symbol(temp, rank=(tensor_index, )), ast.Symbol(builder.coefficient_map[acting_coefficient], rank=(isym, jsym)))) # Outer-loop running over node_extent loop = ast.For(ast.Decl("unsigned int", isym, init=0), ast.Less(isym, node_extent), ast.Incr(isym, 1), inner_loop) aux_statements.append(loop) aux_temps[acting_coefficient] = temp else: raise NotImplementedError( "Auxiliary expression type %s not currently implemented." % type(exp)) return aux_temps, aux_statements
def get_restriction_kernel(fiat_element, unique_indices, dim=1, no_weights=False): weights = restriction_weights(fiat_element)[unique_indices].T ncdof = weights.shape[0] nfdof = weights.shape[1] arglist = [ast.Decl("double", ast.Symbol("coarse", (ncdof*dim, ))), ast.Decl("double *restrict *restrict ", ast.Symbol("fine", ()), qualifiers=["const"])] if not no_weights: arglist.append(ast.Decl("double *restrict *restrict", ast.Symbol("count_weights", ()), qualifiers=["const"])) all_ones = np.allclose(weights, 1.0) if all_ones: w = [] else: w_sym = ast.Symbol("weights", (ncdof, nfdof)) init = ast.ArrayInit(format_array_literal(weights)) w = [ast.Decl("double", w_sym, init, qualifiers=["const"])] i = ast.Symbol("i", ()) j = ast.Symbol("j", ()) k = ast.Symbol("k", ()) fine = ast.Symbol("fine", (j, k)) if no_weights: if all_ones: assign = fine else: assign = ast.Prod(fine, ast.Symbol("weights", (i, j))) else: if all_ones: assign = ast.Prod(fine, ast.Symbol("count_weights", (j, 0))) else: assign = ast.Prod(fine, ast.Prod(ast.Symbol("weights", (i, j)), ast.Symbol("count_weights", (j, 0)))) assignment = ast.Incr(ast.Symbol("coarse", (ast.Sum(k, ast.Prod(i, ast.c_sym(dim))),)), assign) k_loop = ast.For(ast.Decl("int", k, ast.c_sym(0)), ast.Less(k, ast.c_sym(dim)), ast.Incr(k, ast.c_sym(1)), ast.Block([assignment], open_scope=True)) j_loop = ast.For(ast.Decl("int", j, ast.c_sym(0)), ast.Less(j, ast.c_sym(nfdof)), ast.Incr(j, ast.c_sym(1)), ast.Block([k_loop], open_scope=True)) i_loop = ast.For(ast.Decl("int", i, ast.c_sym(0)), ast.Less(i, ast.c_sym(ncdof)), ast.Incr(i, ast.c_sym(1)), ast.Block([j_loop], open_scope=True)) k = ast.FunDecl("void", "restriction", arglist, ast.Block(w + [i_loop]), pred=["static", "inline"]) return op2.Kernel(k, "restriction", opts=parameters["coffee"])
def get_injection_kernel(fiat_element, unique_indices, dim=1): weights = injection_weights(fiat_element)[unique_indices].T ncdof = weights.shape[0] nfdof = weights.shape[1] # What if we have multiple nodes in same location (DG)? Divide by # rowsum. weights = weights / np.sum(weights, axis=1).reshape(-1, 1) all_same = np.allclose(weights, weights[0, 0]) arglist = [ ast.Decl("double", ast.Symbol("coarse", (ncdof * dim, ))), ast.Decl("double *restrict *restrict ", ast.Symbol("fine", ()), qualifiers=["const"]) ] if all_same: w_sym = ast.Symbol("weights", ()) w = [ast.Decl("double", w_sym, weights[0, 0], qualifiers=["const"])] else: init = ast.ArrayInit(format_array_literal(weights)) w_sym = ast.Symbol("weights", (ncdof, nfdof)) w = [ast.Decl("double", w_sym, init, qualifiers=["const"])] i = ast.Symbol("i", ()) j = ast.Symbol("j", ()) k = ast.Symbol("k", ()) if all_same: assign = ast.Prod(ast.Symbol("fine", (j, k)), w_sym) else: assign = ast.Prod(ast.Symbol("fine", (j, k)), ast.Symbol("weights", (i, j))) assignment = ast.Incr( ast.Symbol("coarse", (ast.Sum(k, ast.Prod(i, ast.c_sym(dim))), )), assign) k_loop = ast.For(ast.Decl("int", k, ast.c_sym(0)), ast.Less(k, ast.c_sym(dim)), ast.Incr(k, ast.c_sym(1)), ast.Block([assignment], open_scope=True)) j_loop = ast.For(ast.Decl("int", j, ast.c_sym(0)), ast.Less(j, ast.c_sym(nfdof)), ast.Incr(j, ast.c_sym(1)), ast.Block([k_loop], open_scope=True)) i_loop = ast.For(ast.Decl("int", i, ast.c_sym(0)), ast.Less(i, ast.c_sym(ncdof)), ast.Incr(i, ast.c_sym(1)), ast.Block([j_loop], open_scope=True)) k = ast.FunDecl("void", "injection", arglist, ast.Block(w + [i_loop]), pred=["static", "inline"]) return op2.Kernel(k, "injection", opts=parameters["coffee"])
def expression_kernel(expr, args): """Produce a :class:`pyop2.Kernel` from the processed UFL expression expr and the corresponding args.""" # Empty slot indicating assignment to indexed LHS, so don't do anything if type(expr) is Zero: return fs = args[0].function.function_space() d = ast.Symbol("dim") ast_expr = _ast(expr) body = ast.Block( ( ast.Decl("int", d), ast.For(ast.Assign(d, ast.Symbol(0)), ast.Less(d, ast.Symbol(fs.dof_dset.cdim)), ast.Incr(d, ast.Symbol(1)), ast_expr) ) ) return op2.Kernel(ast.FunDecl("void", "expression", [arg.arg for arg in args], body), "expression")
def get_prolongation_kernel(fiat_element, unique_indices, dim=1): weights = get_restriction_weights(fiat_element)[unique_indices] nfdof = weights.shape[0] ncdof = weights.shape[1] arglist = [ ast.Decl("double", ast.Symbol("fine", (nfdof * dim, ))), ast.Decl("double", ast.Symbol("*restrict *restrict coarse", ()), qualifiers=["const"]) ] all_same = np.allclose(weights, weights[0, 0]) if all_same: w_sym = ast.Symbol("weights", ()) w = [ast.Decl("double", w_sym, weights[0, 0], qualifiers=["const"])] else: w_sym = ast.Symbol("weights", (nfdof, ncdof)) init = ast.ArrayInit(format_array_literal(weights)) w = [ast.Decl("double", w_sym, init, qualifiers=["const"])] i = ast.Symbol("i", ()) j = ast.Symbol("j", ()) k = ast.Symbol("k", ()) if all_same: assign = ast.Prod(ast.Symbol("coarse", (j, k)), w_sym) else: assign = ast.Prod(ast.Symbol("coarse", (j, k)), ast.Symbol("weights", (i, j))) assignment = ast.Incr( ast.Symbol("fine", (ast.Sum(k, ast.Prod(i, ast.c_sym(dim))), )), assign) k_loop = ast.For(ast.Decl("int", k, ast.c_sym(0)), ast.Less(k, ast.c_sym(dim)), ast.Incr(k, ast.c_sym(1)), ast.Block([assignment], open_scope=True)) j_loop = ast.For(ast.Decl("int", j, ast.c_sym(0)), ast.Less(j, ast.c_sym(ncdof)), ast.Incr(j, ast.c_sym(1)), ast.Block([k_loop], open_scope=True)) i_loop = ast.For(ast.Decl("int", i, ast.c_sym(0)), ast.Less(i, ast.c_sym(nfdof)), ast.Incr(i, ast.c_sym(1)), ast.Block([j_loop], open_scope=True)) k = ast.FunDecl("void", "prolongation", arglist, ast.Block(w + [i_loop]), pred=["static", "inline"]) return op2.Kernel(k, "prolongation", opts=parameters["coffee"])
def statement_for(tree, parameters): extent = tree.index.extent assert extent i = _coffee_symbol(parameters.index_names[tree.index]) # TODO: symbolic constant for "int" return coffee.For(coffee.Decl("int", i, init=0), coffee.Less(i, extent), coffee.Incr(i, 1), statement(tree.children[0], parameters))
def get_count_kernel(arity): arglist = [ast.Decl("double", ast.Symbol("weight", (arity, )))] i = ast.Symbol("i", ()) assignment = ast.Incr(ast.Symbol("weight", (i, )), ast.c_sym(1.0)) loop = ast.For(ast.Decl("int", i, ast.c_sym(0)), ast.Less(i, ast.c_sym(arity)), ast.Incr(i, ast.c_sym(1)), ast.Block([assignment], open_scope=True)) k = ast.FunDecl("void", "count_weights", arglist, ast.Block([loop]), pred=["static", "inline"]) return op2.Kernel(k, "count_weights", opts=parameters["coffee"])
def statement_for(tree, parameters): if tree.index in parameters.argument_indices: pragma = "#pragma coffee linear loop" else: pragma = None extent = tree.index.extent assert extent i = _coffee_symbol(parameters.index_names[tree.index]) # TODO: symbolic constant for "int" return coffee.For(coffee.Decl("int", i, init=0), coffee.Less(i, extent), coffee.Incr(i, 1), statement(tree.children[0], parameters), pragma=pragma)
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 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
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), )
def exterior_facet_boundary_node_map(self, method): '''The :class:`pyop2.Map` from exterior facets to the nodes on those facets. Note that this differs from :meth:`exterior_facet_node_map` in that only surface nodes are referenced, not all nodes in cells touching the surface. :arg method: The method for determining boundary nodes. See :class:`~.bcs.DirichletBC`. ''' el = self.fiat_element dim = self._mesh.facet_dimension() if method == "topological": boundary_dofs = el.entity_closure_dofs()[dim] elif method == "geometric": boundary_dofs = el.facet_support_dofs() nodes_per_facet = \ len(boundary_dofs[0]) # HACK ALERT # The facet set does not have a halo associated with it, since # we only construct halos for DoF sets. Fortunately, this # loop is direct and we already have all the correct # information available locally. So We fake a set of the # correct size and carry out a direct loop facet_set = op2.Set(self._mesh.exterior_facets.set.total_size) fs_dat = op2.Dat(facet_set**el.space_dimension(), data=self.exterior_facet_node_map().values_with_halo) facet_dat = op2.Dat(facet_set**nodes_per_facet, dtype=np.int32) local_facet_nodes = np.array( [dofs for e, dofs in boundary_dofs.iteritems()]) # Helper function to turn the inner index of an array into c # array literals. c_array = lambda xs: "{" + ", ".join(map(str, xs)) + "}" body = ast.Block([ ast.Decl("int", ast.Symbol("l_nodes", (len(el.get_reference_element().topology[dim]), nodes_per_facet)), init=ast.ArrayInit( c_array(map(c_array, local_facet_nodes))), qualifiers=["const"]), ast.For( ast.Decl("int", "n", 0), ast.Less("n", nodes_per_facet), ast.Incr("n", 1), ast.Assign( ast.Symbol("facet_nodes", ("n", )), ast.Symbol("cell_nodes", ("l_nodes[facet[0]][n]", )))) ]) kernel = op2.Kernel( ast.FunDecl("void", "create_bc_node_map", [ ast.Decl("int*", "cell_nodes"), ast.Decl("int*", "facet_nodes"), ast.Decl("unsigned int*", "facet") ], body), "create_bc_node_map") local_facet_dat = op2.Dat( facet_set**self._mesh.exterior_facets._rank, self._mesh.exterior_facets.local_facet_dat.data_ro_with_halos, dtype=np.uintc) op2.par_loop(kernel, facet_set, fs_dat(op2.READ), facet_dat(op2.WRITE), local_facet_dat(op2.READ)) if isinstance(self._mesh, mesh_t.ExtrudedMesh): offset = self.offset[boundary_dofs[0]] else: offset = None return op2.Map(facet_set, self.node_set, nodes_per_facet, facet_dat.data_ro_with_halos, name="exterior_facet_boundary_node", offset=offset)
def compile_c_kernel(expression, to_pts, to_element, fs, coords): """Produce a :class:`PyOP2.Kernel` from the c expression provided.""" coords_space = coords.function_space() coords_element = create_element(coords_space.ufl_element(), vector_is_mixed=False) names = {v[0] for v in expression._user_args} X = list(coords_element.tabulate(0, to_pts).values())[0] # Produce C array notation of X. X_str = "{{" + "},\n{".join([",".join(map(str, x)) for x in X.T]) + "}}" A = utils.unique_name("A", names) X = utils.unique_name("X", names) x_ = utils.unique_name("x_", names) k = utils.unique_name("k", names) d = utils.unique_name("d", names) i_ = utils.unique_name("i", names) # x is a reserved name. x = "x" if "x" in names: raise ValueError( "cannot use 'x' as a user-defined Expression variable") ass_exp = [ ast.Assign(ast.Symbol(A, (k, ), ((len(expression.code), i), )), ast.FlatBlock("%s" % code)) for i, code in enumerate(expression.code) ] dim = coords_space.value_size ndof = to_element.space_dimension() xndof = coords_element.space_dimension() nfdof = to_element.space_dimension() * numpy.prod(fs.value_size, dtype=int) init_X = ast.Decl(typ="double", sym=ast.Symbol(X, rank=(ndof, xndof)), qualifiers=["const"], init=X_str) init_x = ast.Decl(typ="double", sym=ast.Symbol(x, rank=(coords_space.value_size, ))) init_pi = ast.Decl(typ="double", sym="pi", qualifiers=["const"], init="3.141592653589793") init = ast.Block([init_X, init_x, init_pi]) incr_x = ast.Incr( ast.Symbol(x, rank=(d, )), ast.Prod(ast.Symbol(X, rank=(k, i_)), ast.Symbol(x_, rank=(ast.Sum(ast.Prod(i_, dim), d), )))) assign_x = ast.Assign(ast.Symbol(x, rank=(d, )), 0) loop_x = ast.For(init=ast.Decl("unsigned int", i_, 0), cond=ast.Less(i_, xndof), incr=ast.Incr(i_, 1), body=[incr_x]) block = ast.For(init=ast.Decl("unsigned int", d, 0), cond=ast.Less(d, dim), incr=ast.Incr(d, 1), body=[assign_x, loop_x]) loop = ast.c_for(k, ndof, ast.Block([block] + ass_exp, open_scope=True)) user_args = [] user_init = [] for _, arg in expression._user_args: if arg.shape == (1, ): user_args.append(ast.Decl("double *", "%s_" % arg.name)) user_init.append( ast.FlatBlock("const double %s = *%s_;" % (arg.name, arg.name))) else: user_args.append(ast.Decl("double *", arg.name)) kernel_code = ast.FunDecl( "void", "expression_kernel", [ ast.Decl("double", ast.Symbol(A, (nfdof, ))), ast.Decl("double*", x_) ] + user_args, ast.Block(user_init + [init, loop], open_scope=False)) coefficients = [coords] for _, arg in expression._user_args: coefficients.append(GlobalWrapper(arg)) return op2.Kernel(kernel_code, kernel_code.name), False, tuple(coefficients)
def init_cell_orientations(self, expr): """Compute and initialise :attr:`cell_orientations` relative to a specified orientation. :arg expr: an :class:`.Expression` evaluated to produce a reference normal direction. """ import firedrake.function as function import firedrake.functionspace as functionspace if expr.value_shape()[0] != 3: raise NotImplementedError('Only implemented for 3-vectors') if self.ufl_cell() not in (ufl.Cell('triangle', 3), ufl.Cell("quadrilateral", 3), ufl.OuterProductCell(ufl.Cell('interval'), ufl.Cell('interval'), gdim=3)): raise NotImplementedError('Only implemented for triangles and quadrilaterals embedded in 3d') if hasattr(self.topology, '_cell_orientations'): raise RuntimeError("init_cell_orientations already called, did you mean to do so again?") v0 = lambda x: ast.Symbol("v0", (x,)) v1 = lambda x: ast.Symbol("v1", (x,)) n = lambda x: ast.Symbol("n", (x,)) x = lambda x: ast.Symbol("x", (x,)) coords = lambda x, y: ast.Symbol("coords", (x, y)) body = [] body += [ast.Decl("double", v(3)) for v in [v0, v1, n, x]] body.append(ast.Decl("double", "dot")) body.append(ast.Assign("dot", 0.0)) body.append(ast.Decl("int", "i")) # if triangle, use v0 = x1 - x0, v1 = x2 - x0 # otherwise, for the various quads, use v0 = x2 - x0, v1 = x1 - x0 # recall reference element ordering: # triangle: 2 quad: 1 3 # 0 1 0 2 if self.ufl_cell() == ufl.Cell('triangle', 3): body.append(ast.For(ast.Assign("i", 0), ast.Less("i", 3), ast.Incr("i", 1), [ast.Assign(v0("i"), ast.Sub(coords(1, "i"), coords(0, "i"))), ast.Assign(v1("i"), ast.Sub(coords(2, "i"), coords(0, "i"))), ast.Assign(x("i"), 0.0)])) else: body.append(ast.For(ast.Assign("i", 0), ast.Less("i", 3), ast.Incr("i", 1), [ast.Assign(v0("i"), ast.Sub(coords(2, "i"), coords(0, "i"))), ast.Assign(v1("i"), ast.Sub(coords(1, "i"), coords(0, "i"))), ast.Assign(x("i"), 0.0)])) # n = v0 x v1 body.append(ast.Assign(n(0), ast.Sub(ast.Prod(v0(1), v1(2)), ast.Prod(v0(2), v1(1))))) body.append(ast.Assign(n(1), ast.Sub(ast.Prod(v0(2), v1(0)), ast.Prod(v0(0), v1(2))))) body.append(ast.Assign(n(2), ast.Sub(ast.Prod(v0(0), v1(1)), ast.Prod(v0(1), v1(0))))) body.append(ast.For(ast.Assign("i", 0), ast.Less("i", 3), ast.Incr("i", 1), [ast.Incr(x(j), coords("i", j)) for j in range(3)])) body.extend([ast.FlatBlock("dot += (%(x)s) * n[%(i)d];\n" % {"x": x_, "i": i}) for i, x_ in enumerate(expr.code)]) body.append(ast.Assign("orientation[0][0]", ast.Ternary(ast.Less("dot", 0), 1, 0))) kernel = op2.Kernel(ast.FunDecl("void", "cell_orientations", [ast.Decl("int**", "orientation"), ast.Decl("double**", "coords")], ast.Block(body)), "cell_orientations") # Build the cell orientations as a DG0 field (so that we can # pass it in for facet integrals and the like) fs = functionspace.FunctionSpace(self, 'DG', 0) cell_orientations = function.Function(fs, name="cell_orientations", dtype=np.int32) op2.par_loop(kernel, self.cell_set, cell_orientations.dat(op2.WRITE, cell_orientations.cell_node_map()), self.coordinates.dat(op2.READ, self.coordinates.cell_node_map())) self.topology._cell_orientations = cell_orientations
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 _generate_integral_ir(points, terms, sets, optimise_parameters, parameters): "Generate code to evaluate the element tensor." # For checking if the integral code is for a matrix def is_matrix(loop): loop_indices = [ l[0] for l in loop ] return (format["first free index"] in loop_indices and \ format["second free index"] in loop_indices) # Prefetch formats to speed up code generation. p_format = parameters["format"] f_comment = format["comment"] f_mul = format["mul"] f_scale_factor = format["scale factor"] f_iadd = format["iadd"] f_add = format["add"] f_A = format["element tensor"][p_format] f_j = format["first free index"] f_k = format["second free index"] f_loop = format["generate loop"] f_B = format["basis constant"] # Initialise return values. code = [] num_ops = 0 loops = {} # Extract sets. used_weights, used_psi_tables, used_nzcs, trans_set = sets nests = [] # Loop terms and create code. for loop, (data, entry_vals) in terms.items(): # If we don't have any entry values, there's no need to generate the loop. if not entry_vals: continue # Get data. t_set, u_weights, u_psi_tables, u_nzcs, basis_consts = data # If we have a value, then we also need to update the sets of used variables. trans_set.update(t_set) used_weights.update(u_weights) used_psi_tables.update(u_psi_tables) used_nzcs.update(u_nzcs) # @@@: A[0][0] += FE0[ip][j]*FE0[ip][k]*W24[ip]*det; entry_ir = [] for entry, value, ops in entry_vals: # Left hand side it_vars = entry if len(loop) > 0 else (0,) rank = tuple(i.loop_index if hasattr(i, 'loop_index') else i for i in it_vars) offset = tuple((1, int(i.offset)) if hasattr(i, 'offset') else (1, 0) for i in it_vars) local_tensor = pyop2.Symbol(f_A(''), rank, offset) # Right hand side pyop2_rhs = visit_rhs(value) entry_ir.append(pyop2.Incr(local_tensor, pyop2_rhs, "#pragma coffee expression")) # Disgusting hack, ensure resulting IR comes out in same order # on all processes. entry_ir = sorted(entry_ir, key=lambda x: x.gencode()) if len(loop) == 0: nest = pyop2.Block(entry_ir, open_scope=True) elif len(loop) in [1, 2]: it_var = c_sym(loop[0][0]) end = c_sym(loop[0][2]) nest = pyop2.For(pyop2.Decl("int", it_var, c_sym(0)), pyop2.Less(it_var, end), \ pyop2.Incr(it_var, c_sym(1)), pyop2.Block(entry_ir, open_scope=True), "#pragma coffee itspace") if len(loop) == 2: it_var = c_sym(loop[1][0]) end = c_sym(loop[1][2]) nest_k = pyop2.For(pyop2.Decl("int", it_var, c_sym(0)), pyop2.Less(it_var, end), \ pyop2.Incr(it_var, c_sym(1)), pyop2.Block(entry_ir, open_scope=True), "#pragma coffee itspace") nest.children[0] = pyop2.Block([nest_k], open_scope=True) nests.append(nest) return nests, num_ops
def _generate_element_tensor(integrals, sets, optimise_parameters, parameters): "Construct quadrature code for element tensors." # Prefetch formats to speed up code generation. f_comment = format["comment"] f_ip = format["integration points"] f_I = format["ip constant"] f_loop = format["generate loop"] f_ip_coords = format["generate ip coordinates"] f_coords = format["coordinate_dofs"] f_double = format["float declaration"] f_decl = format["declaration"] f_X = format["ip coordinates"] f_C = format["conditional"] # Initialise return values. tensor_ops_count = 0 ffc_assert(1 == len(integrals), "This function is not capable of handling multiple integrals.") # We receive a dictionary {num_points: form,}. # Loop points and forms. for points, terms, functions, ip_consts, coordinate, conditionals in integrals: nest_ir = [] ip_ir = [] num_ops = 0 # Generate code to compute coordinates if used. if coordinate: raise RuntimeError("Don't know how to compute coordinates") # Left in place for posterity name, gdim, ip, r = coordinate element_code += ["", f_comment("Declare array to hold physical coordinate of quadrature point.")] element_code += [f_decl(f_double, f_X(points, gdim))] ops, coord_code = f_ip_coords(gdim, points, name, ip, r) ip_code += ["", f_comment("Compute physical coordinate of quadrature point, operations: %d." % ops)] ip_code += [coord_code] num_ops += ops # Update used psi tables and transformation set. sets[1].add(name) sets[3].add(f_coords(r)) # Generate code to compute function values. if functions: const_func_code, func_code, ops = _generate_functions(functions, sets) nest_ir += const_func_code ip_ir += func_code num_ops += ops # Generate code to compute conditionals (might depend on coordinates # and function values so put here). # TODO: Some conditionals might only depend on geometry so they # should be moved outside if possible. if conditionals: ip_ir.append(pyop2.Decl(f_double, c_sym(f_C(len(conditionals))))) # Sort conditionals (need to in case of nested conditionals). reversed_conds = dict([(n, (o, e)) for e, (t, o, n) in conditionals.items()]) for num in range(len(conditionals)): name = format["conditional"](num) ops, expr = reversed_conds[num] ip_ir.append(pyop2.Assign(c_sym(name), visit_rhs(expr))) num_ops += ops # Generate code for ip constant declarations. # TODO: this code should be removable as only executed when ffc's optimisations are on ip_const_ops, ip_const_code = generate_aux_constants(ip_consts, f_I,\ format["assign"], True) if len(ip_const_code) > 0: raise RuntimeError("IP Const code not supported") num_ops += ip_const_ops # Generate code to evaluate the element tensor. code, ops = _generate_integral_ir(points, terms, sets, optimise_parameters, parameters) num_ops += ops tensor_ops_count += num_ops*points ip_ir += code # Loop code over all IPs. # @@@: for (ip ...) { A[0][0] += ... } if points > 1: it_var = pyop2.Symbol(f_ip, ()) nest_ir += [pyop2.For(pyop2.Decl("int", it_var, c_sym(0)), pyop2.Less(it_var, c_sym(points)), pyop2.Incr(it_var, c_sym(1)), pyop2.Block(ip_ir, open_scope=True))] else: nest_ir += ip_ir return (nest_ir, tensor_ops_count)
def _generate_functions(functions, sets): "Generate declarations for functions and code to compute values." f_comment = format["comment"] f_double = format["float declaration"] f_F = format["function value"] f_float = format["floating point"] f_decl = format["declaration"] f_r = format["free indices"] f_iadd = format["iadd"] f_loop = format["generate loop"] # Create the function declarations -- only the (unique) variables we need const_vardecls = set(fd.id for fd in functions.values() if fd.cellwise_constant) const_ast_items = [pyop2.Decl(f_double, c_sym(f_F(n)), c_sym(f_float(0))) for n in const_vardecls] vardecls = set(fd.id for fd in functions.values() if not fd.cellwise_constant) ast_items = [pyop2.Decl(f_double, c_sym(f_F(n)), c_sym(f_float(0))) for n in vardecls] # Get sets. used_psi_tables = sets[1] used_nzcs = sets[2] # Sort functions after being cellwise constant and loop ranges. function_groups = collections.defaultdict(list) for f, fd in functions.items(): function_groups[(fd.cellwise_constant, fd.loop_range)].append(f) total_ops = 0 # Loop ranges and get list of functions. for (cellwise_constant, loop_range), function_list in function_groups.iteritems(): function_expr = [] # Loop functions. func_ops = 0 for function in function_list: data = functions[function] range_i = data.loop_range if not isinstance(range_i, tuple): range_i = tuple([range_i]) # Add name to used psi names and non zeros name to used_nzcs. used_psi_tables.add(data.psi_name) used_nzcs.update(data.used_nzcs) # # TODO: This check can be removed for speed later. # REMOVED this, since we might need to increment into the same # number more than once for mixed element + interior facets # ffc_assert(data.id not in function_expr, "This is definitely not supposed to happen!") # Convert function to COFFEE ast node, save string # representation for sorting (such that we're reproducible # in parallel). function = visit_rhs(function) key = str(function) function_expr.append((data.id, function, key)) # Get number of operations to compute entry and add to function operations count. func_ops += (data.ops + 1)*sum(range_i) # Gather, sorted by string rep of function. lines = [pyop2.Incr(c_sym(f_F(n)), fn) for n, fn, _ in sorted(function_expr, key=lambda x: x[2])] if isinstance(loop_range, tuple): if not all(map(lambda x: x==loop_range[0], loop_range)): raise RuntimeError("General mixed elements not yet supported in PyOP2") loop_vars = [ (f_r[0], 0, loop_range[0]), (f_r[1], 0, len(loop_range)) ] else: loop_vars = [(f_r[0], 0, loop_range)] # TODO: If loop_range == 1, this loop may be unneccessary. Not sure if it's safe to just skip it. it_var = c_sym(loop_vars[0][0]) loop_size = c_sym(loop_vars[0][2]) ast_item = pyop2.For(pyop2.Decl("int", it_var, c_sym(0)), pyop2.Less(it_var, loop_size), pyop2.Incr(it_var, c_sym(1)), pyop2.Block(lines, open_scope=True)) if cellwise_constant: const_ast_items.append(ast_item) else: ast_items.append(ast_item) return const_ast_items, ast_items, total_ops
def auxiliary_temporaries(builder, declared_temps): """This function generates auxiliary information regarding special handling of expressions that require creating additional temporaries. :arg builder: a :class:`KernelBuilder` object that contains all the necessary temporary and expression information. :arg declared_temps: a `dict` of temporaries that have already been declared and assigned values. This will be updated in this method and referenced later in the compiler. Returns: a list of auxiliary statements are returned that contain temporary declarations and any code-blocks needed to evaluate the expression. """ aux_statements = [] for exp in builder.aux_exprs: if isinstance(exp, Inverse): if builder._ref_counts[exp] > 1: # Get the temporary for the particular expression result = metaphrase_slate_to_cpp(exp, declared_temps) # Now we use the generated result and assign the value to the # corresponding temporary. temp = ast.Symbol("auxT%d" % len(declared_temps)) shape = exp.shape aux_statements.append( ast.Decl(eigen_matrixbase_type(shape), temp)) aux_statements.append(ast.FlatBlock("%s.setZero();\n" % temp)) aux_statements.append(ast.Assign(temp, result)) # Update declared temps declared_temps[exp] = temp elif isinstance(exp, Action): # Action computations are relatively inexpensive, so # we don't waste memory space on creating temps for # these expressions. However, we must create a temporary # for the actee coefficient (if we haven't already). actee, = exp.actee if actee not in declared_temps: # Declare a temporary for the coefficient V = actee.function_space() shape_array = [(Vi.finat_element.space_dimension(), np.prod(Vi.shape)) for Vi in V.split()] ctemp = ast.Symbol("auxT%d" % len(declared_temps)) shape = sum(n * d for (n, d) in shape_array) typ = eigen_matrixbase_type(shape=(shape, )) aux_statements.append(ast.Decl(typ, ctemp)) aux_statements.append(ast.FlatBlock("%s.setZero();\n" % ctemp)) # Now we populate the temporary with the coefficient # information and insert in the right place. offset = 0 for i, shp in enumerate(shape_array): node_extent, dof_extent = shp # Now we unpack the function and insert its entries into a # 1D vector temporary isym = ast.Symbol("i1") jsym = ast.Symbol("j1") tensor_index = ast.Sum( offset, ast.Sum(ast.Prod(dof_extent, isym), jsym)) # Inner-loop running over dof_extent coeff_sym = ast.Symbol(builder.coefficient(actee)[i], rank=(isym, jsym)) coeff_temp = ast.Symbol(ctemp, rank=(tensor_index, )) inner_loop = ast.For( ast.Decl("unsigned int", jsym, init=0), ast.Less(jsym, dof_extent), ast.Incr(jsym, 1), ast.Assign(coeff_temp, coeff_sym)) # Outer-loop running over node_extent loop = ast.For(ast.Decl("unsigned int", isym, init=0), ast.Less(isym, node_extent), ast.Incr(isym, 1), inner_loop) aux_statements.append(loop) offset += node_extent * dof_extent # Update declared temporaries with the coefficient declared_temps[actee] = ctemp else: raise NotImplementedError( "Auxiliary expr type %s not currently implemented." % type(exp)) return aux_statements
def coefficient_temporaries(builder, declared_temps): """Generates coefficient temporary statements for assigning coefficients to vector temporaries. :arg builder: The :class:`LocalKernelBuilder` containing all relevant expression information. :arg declared_temps: A `dict` keeping track of all declared temporaries. This dictionary is updated as coefficients are assigned temporaries. 'AssembledVector's require creating coefficient temporaries to store data. The temporaries are created by inspecting the function space of the coefficient to compute node and dof extents. The coefficient is then assigned values by looping over both the node extent and dof extent (double FOR-loop). A double FOR-loop is needed for each function space (if the function space is mixed, then a loop will be constructed for each component space). The general structure of each coefficient loop will be: FOR (i1=0; i1<node_extent; i1++): FOR (j1=0; j1<dof_extent; j1++): VT0[offset + (dof_extent * i1) + j1] = w_0_0[i1][j1] VT1[offset + (dof_extent * i1) + j1] = w_1_0[i1][j1] . . . where wT0, wT1, ... are temporaries for coefficients sharing the same node and dof extents. The offset is computed based on whether the function space is mixed. The offset is always 0 for non-mixed coefficients. If the coefficient is mixed, then the offset is incremented by the total number of nodal unknowns associated with the component spaces of the mixed space. """ statements = [ast.FlatBlock("/* Coefficient temporaries */\n")] j = ast.Symbol("j1") loops = [ast.FlatBlock("/* Loops for coefficient temps */\n")] for dofs, cinfo_list in builder.coefficient_vecs.items(): # Collect all coefficients which share the same node/dof extent assignments = [] for cinfo in cinfo_list: fs_i = cinfo.space_index offset = cinfo.offset_index c_shape = cinfo.shape vector = cinfo.vector function = vector._function t = cinfo.local_temp if vector not in declared_temps: # Declare and initialize coefficient temporary c_type = eigen_matrixbase_type(shape=c_shape) statements.append(ast.Decl(c_type, t)) declared_temps[vector] = t # Assigning coefficient values into temporary coeff_sym = ast.Symbol(builder.coefficient(function)[fs_i], rank=(j, )) index = ast.Sum(offset, j) coeff_temp = ast.Symbol(t, rank=(index, )) assignments.append(ast.Assign(coeff_temp, coeff_sym)) # loop over dofs loop = ast.For(ast.Decl("unsigned int", j, init=0), ast.Less(j, dofs), ast.Incr(j, 1), assignments) loops.append(loop) statements.extend(loops) return statements
_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 exterior_facet_boundary_node_map(self, V, method): """Return the :class:`pyop2.Map` from exterior facets to nodes on the boundary. :arg V: The function space. :arg method: The method for determining boundary nodes. See :class:`~.DirichletBC` for details. """ try: return self.map_caches["boundary_node"][method] except KeyError: pass el = V.finat_element dim = self.mesh.facet_dimension() if method == "topological": boundary_dofs = el.entity_closure_dofs()[dim] elif method == "geometric": # This function is only called on extruded meshes when # asking for the nodes that live on the "vertical" # exterior facets. boundary_dofs = entity_support_dofs(el, dim) nodes_per_facet = \ len(boundary_dofs[0]) # HACK ALERT # The facet set does not have a halo associated with it, since # we only construct halos for DoF sets. Fortunately, this # loop is direct and we already have all the correct # information available locally. So We fake a set of the # correct size and carry out a direct loop facet_set = op2.Set(self.mesh.exterior_facets.set.total_size, comm=self.mesh.comm) fs_dat = op2.Dat( facet_set**el.space_dimension(), data=V.exterior_facet_node_map().values_with_halo.view()) facet_dat = op2.Dat(facet_set**nodes_per_facet, dtype=IntType) # Ensure these come out in sorted order. local_facet_nodes = numpy.array( [boundary_dofs[e] for e in sorted(boundary_dofs.keys())]) # Helper function to turn the inner index of an array into c # array literals. c_array = lambda xs: "{" + ", ".join(map(str, xs)) + "}" # AST for: l_nodes[facet[0]][n] rank_ast = ast.Symbol("l_nodes", rank=(ast.Symbol("facet", rank=(0, )), "n")) body = ast.Block([ ast.Decl("int", ast.Symbol("l_nodes", (len(el.cell.topology[dim]), nodes_per_facet)), init=ast.ArrayInit( c_array(map(c_array, local_facet_nodes))), qualifiers=["const"]), ast.For( ast.Decl("int", "n", 0), ast.Less("n", nodes_per_facet), ast.Incr("n", 1), ast.Assign(ast.Symbol("facet_nodes", ("n", )), ast.Symbol("cell_nodes", (rank_ast, )))) ]) kernel = op2.Kernel( ast.FunDecl("void", "create_bc_node_map", [ ast.Decl("%s*" % as_cstr(fs_dat.dtype), "cell_nodes"), ast.Decl("%s*" % as_cstr(facet_dat.dtype), "facet_nodes"), ast.Decl("unsigned int*", "facet") ], body), "create_bc_node_map") local_facet_dat = op2.Dat( facet_set**self.mesh.exterior_facets._rank, self.mesh.exterior_facets.local_facet_dat.data_ro_with_halos, dtype=numpy.uintc) op2.par_loop(kernel, facet_set, fs_dat(op2.READ), facet_dat(op2.WRITE), local_facet_dat(op2.READ)) if self.extruded: offset = self.offset[boundary_dofs[0]] else: offset = None val = op2.Map(facet_set, self.node_set, nodes_per_facet, facet_dat.data_ro_with_halos, name="exterior_facet_boundary_node", offset=offset) self.map_caches["boundary_node"][method] = val return val
def coefficient_temporaries(builder, declared_temps): """Generates coefficient temporary statements for assigning coefficients to vector temporaries. :arg builder: The :class:`LocalKernelBuilder` containing all relevant expression information. :arg declared_temps: A `dict` keeping track of all declared temporaries. This dictionary is updated as coefficients are assigned temporaries. Action computations require creating coefficient temporaries to compute the matrix-vector product. The temporaries are created by inspecting the function space of the coefficient to compute node and dof extents. The coefficient is then assigned values by looping over both the node extent and dof extent (double FOR-loop). A double FOR-loop is needed for each function space (if the function space is mixed, then a loop will be constructed for each component space). The general structure of each coefficient loop will be: FOR (i1=0; i1<node_extent; i1++): FOR (j1=0; j1<dof_extent; j1++): wT0[offset + (dof_extent * i1) + j1] = w_0_0[i1][j1] wT1[offset + (dof_extent * i1) + j1] = w_1_0[i1][j1] . . . where wT0, wT1, ... are temporaries for coefficients sharing the same node and dof extents. The offset is computed based on whether the function space is mixed. The offset is always 0 for non-mixed coefficients. If the coefficient is mixed, then the offset is incremented by the total number of nodal unknowns associated with the component spaces of the mixed space. """ statements = [ast.FlatBlock("/* Coefficient temporaries */\n")] i_sym = ast.Symbol("i1") j_sym = ast.Symbol("j1") loops = [ast.FlatBlock("/* Loops for coefficient temps */\n")] for (nodes, dofs), cinfo_list in builder.action_coefficients.items(): # Collect all coefficients which share the same node/dof extent assignments = [] for cinfo in cinfo_list: fs_i = cinfo.space_index offset = cinfo.offset_index c_shape = cinfo.shape actee = cinfo.coefficient if actee not in declared_temps: # Declare and initialize coefficient temporary c_type = eigen_matrixbase_type(shape=c_shape) t = ast.Symbol("wT%d" % len(declared_temps)) statements.append(ast.Decl(c_type, t)) statements.append(ast.FlatBlock("%s.setZero();\n" % t)) declared_temps[actee] = t # Assigning coefficient values into temporary coeff_sym = ast.Symbol(builder.coefficient(actee)[fs_i], rank=(i_sym, j_sym)) index = ast.Sum(offset, ast.Sum(ast.Prod(dofs, i_sym), j_sym)) coeff_temp = ast.Symbol(t, rank=(index, )) assignments.append(ast.Assign(coeff_temp, coeff_sym)) # Inner-loop running over dof extent inner_loop = ast.For(ast.Decl("unsigned int", j_sym, init=0), ast.Less(j_sym, dofs), ast.Incr(j_sym, 1), assignments) # Outer-loop running over node extent loop = ast.For(ast.Decl("unsigned int", i_sym, init=0), ast.Less(i_sym, nodes), ast.Incr(i_sym, 1), inner_loop) loops.append(loop) statements.extend(loops) return statements