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
