Ejemplo n.º 1
0
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"])
Ejemplo n.º 2
0
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"])
Ejemplo n.º 3
0
    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
Ejemplo n.º 4
0
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))
Ejemplo n.º 5
0
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"])
Ejemplo n.º 6
0
 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)
Ejemplo n.º 7
0
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
Ejemplo n.º 8
0
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))
Ejemplo n.º 9
0
def test_prod_div():
    tree = ast.Prod("a", ast.Div("1", "b"))

    assert tree.gencode() == "(a) * (((1) / (b)))"
Ejemplo n.º 10
0
    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
Ejemplo n.º 11
0
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
Ejemplo n.º 12
0
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)
Ejemplo n.º 13
0
    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)),
Ejemplo n.º 14
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
Ejemplo n.º 15
0
def test_prod_div():
    tree = ast.Prod("a", ast.Div("1", "b"))

    assert tree.gencode() == "a * (1 / b)"
Ejemplo n.º 16
0
def _expression_product(expr, parameters):
    return coffee.Prod(*[expression(c, parameters)
                         for c in expr.children])
Ejemplo n.º 17
0
    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