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 ast_matmul(self, F_a): """Generate an AST for a PyOP2 kernel performing a matrix-vector multiplication. :param F_a: Assembled firedrake.Function object for the RHS""" # The number of dofs on each element is /ndofs*cdim/ F_a_fs = F_a.function_space() ndofs = sum(F_a_fs.topological.dofs_per_entity) cdim = F_a_fs.dim name = 'mat_vec_mul_kernel_%s' % F_a_fs.name identifier = (ndofs, cdim, name) if identifier in self.asts: return self.asts[identifier] # Craft the AST body = ast.Incr(ast.Symbol('C', ('i/%d' % cdim, 'i%%%d' % cdim)), ast.Prod(ast.Symbol('A', ('i',), ((ndofs*cdim, 'j*%d + k' % cdim),)), ast.Symbol('B', ('j', 'k')))) body = ast.c_for('k', cdim, body).children[0] body = [ast.Assign(ast.Symbol('C', ('i/%d' % cdim, 'i%%%d' % cdim)), '0.0'), ast.c_for('j', ndofs, body).children[0]] body = ast.Root([ast.c_for('i', ndofs*cdim, body).children[0]]) funargs = [ast.Decl('double*', 'A'), ast.Decl('double**', 'B'), ast.Decl('double**', 'C')] fundecl = ast.FunDecl('void', name, funargs, body, ['static', 'inline']) # Track the AST for later fast retrieval self.asts[identifier] = fundecl return fundecl
def _expression_flexiblyindexed(expr, parameters): var = expression(expr.children[0], parameters) assert isinstance(var, coffee.Symbol) assert not var.rank assert not var.offset rank = [] offset = [] for off, idxs in expr.dim2idxs: for index, stride in idxs: assert isinstance(index, gem.Index) if len(idxs) == 0: rank.append(off) offset.append((1, 0)) elif len(idxs) == 1: (index, stride), = idxs rank.append(parameters.index_names[index]) offset.append((stride, off)) else: parts = [] if off: parts += [coffee.Symbol(str(off))] for index, stride in idxs: index_sym = coffee.Symbol(parameters.index_names[index]) assert stride if stride == 1: parts += [index_sym] else: parts += [coffee.Prod(index_sym, coffee.Symbol(str(stride)))] assert parts rank.append(reduce(coffee.Sum, parts)) offset.append((1, 0)) return coffee.Symbol(var.symbol, rank=tuple(rank), offset=tuple(offset))
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 create_pyop2_node(typ, exp1, exp2): """Create an expr node starting from two FFC symbols.""" if typ == 2: return pyop2.Prod(exp1, exp2) if typ == 3: return pyop2.Sum(exp1, exp2) if typ == 4: return pyop2.Div(exp1, exp2)
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 ker_reduce_ind_read(): body = ast.Incr('a', ast.Prod(ast.Symbol('V', (0, 'i')), ast.Symbol('B', (0,)))) body = \ [ast.Decl('int', 'a', '0')] +\ ItSpace().to_for([(0, 2)], ('i',), [body]) +\ [ast.Incr(ast.Symbol('A', (0,)), 'a')] return ast.FunDecl('void', 'ker_reduce_ind_read', [ast.Decl('int', 'A', qualifiers=['unsigned'], pointers=['']), ast.Decl('int', 'V', qualifiers=['unsigned'], pointers=['', '']), ast.Decl('int', 'B', qualifiers=['unsigned'], pointers=[''])], ast.Block(body))
def test_prod_div(): tree = ast.Prod("a", ast.Div("1", "b")) assert tree.gencode() == "(a) * (((1) / (b)))"
def ast_matmul(self, F_a, implementation='optimized'): """Generate an AST for a PyOP2 kernel performing a matrix-vector multiplication.""" # The number of dofs on each element is /ndofs*cdim/ F_a_fs = F_a.function_space() ndofs = F_a_fs.fiat_element.entity_dofs() ndofs = sum(self.mesh.make_dofs_per_plex_entity(ndofs)) cdim = F_a_fs.dim name = 'mat_vec_mul_kernel_%s' % F_a_fs.name identifier = (ndofs, cdim, name, implementation) if identifier in self.asts: return self.asts[identifier] from coffee import isa, options if cdim and cdim % isa['dp_reg'] == 0: simd_pragma = '#pragma simd reduction(+:sum)' else: simd_pragma = '' # Craft the AST if implementation == 'optimized' and cdim >= 4: body = ast.Incr( ast.Symbol('sum'), ast.Prod( ast.Symbol('A', ('i', ), ((ndofs * cdim, 'j*%d + k' % cdim), )), ast.Symbol('B', ('j', 'k')))) body = ast.c_for('k', cdim, body, simd_pragma).children[0] body = [ ast.Decl('const int', ast.Symbol('index'), init=ast.Symbol('i%%%d' % cdim)), ast.Decl('double', ast.Symbol('sum'), init=ast.Symbol('0.0')), ast.c_for('j', ndofs, body).children[0], ast.Assign(ast.Symbol('C', ('i/%d' % cdim, 'index')), 'sum') ] body = ast.Block([ast.c_for('i', ndofs * cdim, body).children[0]]) funargs = [ ast.Decl('double* restrict', 'A'), ast.Decl('double *restrict *restrict', 'B'), ast.Decl('double *restrict *', 'C') ] fundecl = ast.FunDecl('void', name, funargs, body, ['static', 'inline']) else: body = ast.Incr( ast.Symbol('C', ('i/%d' % cdim, 'index')), ast.Prod( ast.Symbol('A', ('i', ), ((ndofs * cdim, 'j*%d + k' % cdim), )), ast.Symbol('B', ('j', 'k')))) body = ast.c_for('k', cdim, body).children[0] body = [ ast.Decl('const int', ast.Symbol('index'), init=ast.Symbol('i%%%d' % cdim)), ast.Assign(ast.Symbol('C', ('i/%d' % cdim, 'index' % cdim)), '0.0'), ast.c_for('j', ndofs, body).children[0] ] body = ast.Block([ast.c_for('i', ndofs * cdim, body).children[0]]) funargs = [ ast.Decl('double* restrict', 'A'), ast.Decl('double *restrict *restrict', 'B'), ast.Decl('double *restrict *', 'C') ] fundecl = ast.FunDecl('void', name, funargs, body, ['static', 'inline']) # Track the AST for later fast retrieval self.asts[identifier] = fundecl return fundecl
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")] 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.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 if vector not in declared_temps: # Declare and initialize coefficient temporary c_type = eigen_matrixbase_type(shape=c_shape) t = ast.Symbol("VT%d" % len(declared_temps)) statements.append(ast.Decl(c_type, t)) statements.append(ast.FlatBlock("%s.setZero();\n" % t)) declared_temps[vector] = t # Assigning coefficient values into temporary coeff_sym = ast.Symbol(builder.coefficient(function)[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
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)
if result._expression_cache is not None: result._expression_cache[key] = vals def assemble_expression(expr, subset=None): """Evaluates UFL expressions on :class:`.Function`\s pointwise and assigns into a new :class:`.Function`.""" result = function.Function(ExpressionWalker().walk(expr)[2]) evaluate_expression(Assign(result, expr), subset) return result _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)),
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 test_prod_div(): tree = ast.Prod("a", ast.Div("1", "b")) assert tree.gencode() == "a * (1 / b)"
def _expression_product(expr, parameters): return coffee.Prod(*[expression(c, parameters) for c in expr.children])
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