def prepare_coordinates(coefficient, name, interior_facet=False):
"""Bridges the kernel interface and the GEM abstraction for
coordinates.
:arg coefficient: UFL Coefficient
:arg name: unique name to refer to the Coefficient in the kernel
:arg interior_facet: interior facet integral?
:returns: (funarg, expression)
funarg - :class:`coffee.Decl` function argument
expression - GEM expression referring to the Coefficient
values
"""
finat_element = create_element(coefficient.ufl_element())
shape = finat_element.index_shape
size = numpy.prod(shape, dtype=int)
if not interior_facet:
funargs = [coffee.Decl(SCALAR_TYPE, coffee.Symbol(name),
pointers=[("",)],
qualifiers=["const"])]
variable = gem.Variable(name, (size,))
expression = gem.reshape(variable, shape)
else:
funargs = [coffee.Decl(SCALAR_TYPE, coffee.Symbol(name+"_0"),
pointers=[("",)],
qualifiers=["const"]),
coffee.Decl(SCALAR_TYPE, coffee.Symbol(name+"_1"),
pointers=[("",)],
qualifiers=["const"])]
variable0 = gem.Variable(name+"_0", (size,))
variable1 = gem.Variable(name+"_1", (size,))
expression = (gem.reshape(variable0, shape),
gem.reshape(variable1, shape))
return funargs, expression
def __init__(self, integral_type, subdomain_id, domain_number):
"""Initialise a kernel builder."""
super(KernelBuilder,
self).__init__(integral_type.startswith("interior_facet"))
self.kernel = Kernel(integral_type=integral_type,
subdomain_id=subdomain_id,
domain_number=domain_number)
self.local_tensor = None
self.coordinates_arg = None
self.coefficient_args = []
self.coefficient_split = {}
# Facet number
if integral_type in ['exterior_facet', 'exterior_facet_vert']:
facet = gem.Variable('facet', (1, ))
self._facet_number = {
None: gem.VariableIndex(gem.Indexed(facet, (0, )))
}
elif integral_type in ['interior_facet', 'interior_facet_vert']:
facet = gem.Variable('facet', (2, ))
self._facet_number = {
'+': gem.VariableIndex(gem.Indexed(facet, (0, ))),
'-': gem.VariableIndex(gem.Indexed(facet, (1, )))
}
elif integral_type == 'interior_facet_horiz':
self._facet_number = {'+': 1, '-': 0}
def test_expressions():
x = gem.Variable("x", (3, 4))
y = gem.Variable("y", (4, ))
i, j = gem.indices(2)
xij = x[i, j]
yj = y[j]
assert xij == gem.Indexed(x, (i, j))
assert yj == gem.Indexed(y, (j, ))
assert xij + yj == gem.Sum(xij, yj)
assert xij * yj == gem.Product(xij, yj)
assert xij - yj == gem.Sum(xij, gem.Product(gem.Literal(-1), yj))
assert xij / yj == gem.Division(xij, yj)
assert xij + 1 == gem.Sum(xij, gem.Literal(1))
assert 1 + xij == gem.Sum(gem.Literal(1), xij)
assert (xij + y).shape == (4, )
assert (x @ y).shape == (3, )
assert x.T.shape == (4, 3)
with pytest.raises(ValueError):
xij.T @ y
with pytest.raises(ValueError):
xij + "foo"
def prepare_arguments(arguments, multiindices, scalar_type, interior_facet=False, diagonal=False):
"""Bridges the kernel interface and the GEM abstraction for
Arguments. Vector Arguments are rearranged here for interior
facet integrals.
:arg arguments: UFL Arguments
:arg multiindices: Argument multiindices
:arg interior_facet: interior facet integral?
:arg diagonal: Are we assembling the diagonal of a rank-2 element tensor?
:returns: (funarg, expression)
funarg - :class:`loopy.GlobalArg` function argument
expressions - GEM expressions referring to the argument
tensor
"""
assert isinstance(interior_facet, bool)
if len(arguments) == 0:
# No arguments
funarg = lp.GlobalArg("A", dtype=scalar_type, shape=(1,))
expression = gem.Indexed(gem.Variable("A", (1,)), (0,))
return funarg, [expression]
elements = tuple(create_element(arg.ufl_element()) for arg in arguments)
shapes = tuple(element.index_shape for element in elements)
if diagonal:
if len(arguments) != 2:
raise ValueError("Diagonal only for 2-forms")
try:
element, = set(elements)
except ValueError:
raise ValueError("Diagonal only for diagonal blocks (test and trial spaces the same)")
elements = (element, )
shapes = tuple(element.index_shape for element in elements)
multiindices = multiindices[:1]
def expression(restricted):
return gem.Indexed(gem.reshape(restricted, *shapes),
tuple(chain(*multiindices)))
u_shape = numpy.array([numpy.prod(shape, dtype=int) for shape in shapes])
if interior_facet:
c_shape = tuple(2 * u_shape)
slicez = [[slice(r * s, (r + 1) * s)
for r, s in zip(restrictions, u_shape)]
for restrictions in product((0, 1), repeat=len(arguments))]
else:
c_shape = tuple(u_shape)
slicez = [[slice(s) for s in u_shape]]
funarg = lp.GlobalArg("A", dtype=scalar_type, shape=c_shape)
varexp = gem.Variable("A", c_shape)
expressions = [expression(gem.view(varexp, *slices)) for slices in slicez]
return funarg, prune(expressions)
def prepare_coefficient(coefficient, name, interior_facet=False):
"""Bridges the kernel interface and the GEM abstraction for
Coefficients.
:arg coefficient: UFL Coefficient
:arg name: unique name to refer to the Coefficient in the kernel
:arg interior_facet: interior facet integral?
:returns: (funarg, expression)
funarg - :class:`coffee.Decl` function argument
expression - GEM expression referring to the Coefficient
values
"""
assert isinstance(interior_facet, bool)
if coefficient.ufl_element().family() == 'Real':
# Constant
funarg = coffee.Decl(SCALAR_TYPE,
coffee.Symbol(name),
pointers=[("restrict", )],
qualifiers=["const"])
expression = gem.reshape(gem.Variable(name, (None, )),
coefficient.ufl_shape)
return funarg, expression
finat_element = create_element(coefficient.ufl_element())
if isinstance(finat_element, TensorFiniteElement):
scalar_shape = finat_element.base_element.index_shape
tensor_shape = finat_element.index_shape[len(scalar_shape):]
else:
scalar_shape = finat_element.index_shape
tensor_shape = ()
scalar_size = numpy.prod(scalar_shape, dtype=int)
tensor_size = numpy.prod(tensor_shape, dtype=int)
funarg = coffee.Decl(SCALAR_TYPE,
coffee.Symbol(name),
pointers=[("const", "restrict"), ("restrict", )],
qualifiers=["const"])
if not interior_facet:
expression = gem.reshape(
gem.Variable(name, (scalar_size, tensor_size)), scalar_shape,
tensor_shape)
else:
varexp = gem.Variable(name, (2 * scalar_size, tensor_size))
plus = gem.view(varexp, slice(scalar_size), slice(tensor_size))
minus = gem.view(varexp, slice(scalar_size, 2 * scalar_size),
slice(tensor_size))
expression = (gem.reshape(plus, scalar_shape, tensor_shape),
gem.reshape(minus, scalar_shape, tensor_shape))
return funarg, expression
def to_reference_coordinates(ufl_coordinate_element):
# Set up UFL form
cell = ufl_coordinate_element.cell()
domain = ufl.Mesh(ufl_coordinate_element)
K = ufl.JacobianInverse(domain)
x = ufl.SpatialCoordinate(domain)
x0_element = ufl.VectorElement("Real", cell, 0)
x0 = ufl.Coefficient(ufl.FunctionSpace(domain, x0_element))
expr = ufl.dot(K, x - x0)
# Translation to GEM
C = ufl_utils.coordinate_coefficient(domain)
expr = ufl_utils.preprocess_expression(expr)
expr = ufl_utils.replace_coordinates(expr, C)
expr = ufl_utils.simplify_abs(expr)
builder = firedrake_interface.KernelBuilderBase()
builder._coefficient(C, "C")
builder._coefficient(x0, "x0")
dim = cell.topological_dimension()
point = gem.Variable('X', (dim, ))
context = tsfc.fem.GemPointContext(
interface=builder,
ufl_cell=cell,
precision=parameters["precision"],
point_indices=(),
point_expr=point,
)
translator = tsfc.fem.Translator(context)
ir = map_expr_dag(translator, expr)
# Unroll result
ir = [gem.Indexed(ir, alpha) for alpha in numpy.ndindex(ir.shape)]
# Unroll IndexSums
max_extent = parameters["unroll_indexsum"]
if max_extent:
def predicate(index):
return index.extent <= max_extent
ir = gem.optimise.unroll_indexsum(ir, predicate=predicate)
# Translate to COFFEE
ir = impero_utils.preprocess_gem(ir)
return_variable = gem.Variable('dX', (dim, ))
assignments = [(gem.Indexed(return_variable, (i, )), e)
for i, e in enumerate(ir)]
impero_c = impero_utils.compile_gem(assignments, ())
body = tsfc.coffee.generate(impero_c, {}, parameters["precision"])
body.open_scope = False
return body
def test_pickle_gem(protocol):
f = gem.VariableIndex(gem.Indexed(gem.Variable('facet', (2, )), (1, )))
q = gem.Index()
r = gem.Index()
_1 = gem.Indexed(gem.Literal(numpy.random.rand(3, 6, 8)), (f, q, r))
_2 = gem.Indexed(
gem.view(gem.Variable('w', (None, None)), slice(8), slice(1)), (r, 0))
expr = gem.ComponentTensor(gem.IndexSum(gem.Product(_1, _2), (r, )), (q, ))
unpickled = pickle.loads(pickle.dumps(expr, protocol))
assert repr(expr) == repr(unpickled)
def prepare_coordinates(coefficient, name, scalar_type, interior_facet=False):
"""Bridges the kernel interface and the GEM abstraction for
coordinates.
:arg coefficient: UFL Coefficient
:arg name: unique name to refer to the Coefficient in the kernel
:arg interior_facet: interior facet integral?
:returns: (funarg, expression)
funarg - :class:`coffee.Decl` function argument
expression - GEM expression referring to the Coefficient
values
"""
finat_element = create_element(coefficient.ufl_element())
shape = finat_element.index_shape
size = numpy.prod(shape, dtype=int)
assert isinstance(finat_element, TensorFiniteElement)
scalar_shape = finat_element.base_element.index_shape
tensor_shape = finat_element._shape
transposed_shape = scalar_shape + tensor_shape
scalar_rank = len(scalar_shape)
def transpose(expr):
indices = tuple(gem.Index(extent=extent) for extent in expr.shape)
transposed_indices = indices[scalar_rank:] + indices[:scalar_rank]
return gem.ComponentTensor(gem.Indexed(expr, indices),
transposed_indices)
if not interior_facet:
funargs = [
coffee.Decl(scalar_type,
coffee.Symbol(name),
pointers=[("", )],
qualifiers=["const"])
]
variable = gem.Variable(name, (size, ))
expression = transpose(gem.reshape(variable, transposed_shape))
else:
funargs = [
coffee.Decl(scalar_type,
coffee.Symbol(name + "_0"),
pointers=[("", )],
qualifiers=["const"]),
coffee.Decl(scalar_type,
coffee.Symbol(name + "_1"),
pointers=[("", )],
qualifiers=["const"])
]
variable0 = gem.Variable(name + "_0", (size, ))
variable1 = gem.Variable(name + "_1", (size, ))
expression = (transpose(gem.reshape(variable0, transposed_shape)),
transpose(gem.reshape(variable1, transposed_shape)))
return funargs, expression
def test_refactorise():
f = gem.Variable('f', (3,))
u = gem.Variable('u', (3,))
v = gem.Variable('v', ())
i = gem.Index()
f_i = gem.Indexed(f, (i,))
u_i = gem.Indexed(u, (i,))
def classify(atomics_set, expression):
if expression in atomics_set:
return ATOMIC
for node in traversal([expression]):
if node in atomics_set:
return COMPOUND
return OTHER
classifier = partial(classify, {u_i, v})
# \sum_i 5*(2*u_i + -1*v)*(u_i + v*f)
expr = gem.IndexSum(
gem.Product(
gem.Literal(5),
gem.Product(
gem.Sum(gem.Product(gem.Literal(2), u_i),
gem.Product(gem.Literal(-1), v)),
gem.Sum(u_i, gem.Product(v, f_i))
)
),
(i,)
)
expected = [
Monomial((i,),
(u_i, u_i),
gem.Literal(10)),
Monomial((i,),
(u_i, v),
gem.Product(gem.Literal(5),
gem.Sum(gem.Product(f_i, gem.Literal(2)),
gem.Literal(-1)))),
Monomial((),
(v, v),
gem.Product(gem.Literal(5),
gem.IndexSum(gem.Product(f_i, gem.Literal(-1)),
(i,)))),
]
actual, = collect_monomials([expr], classifier)
assert expected == list(actual)
def __init__(self, interior_facet=False):
"""Initialise a kernel builder.
:arg interior_facet: kernel accesses two cells
"""
super(KernelBuilderBase, self).__init__(interior_facet=interior_facet)
# Cell orientation
if self.interior_facet:
cell_orientations = gem.Variable("cell_orientations", (2, ))
self._cell_orientations = (gem.Indexed(cell_orientations, (0, )),
gem.Indexed(cell_orientations, (1, )))
else:
cell_orientations = gem.Variable("cell_orientations", (1, ))
self._cell_orientations = (gem.Indexed(cell_orientations, (0, )), )
def prepare_coefficient(coefficient, num, name, interior_facet=False):
"""Bridges the kernel interface and the GEM abstraction for
Coefficients.
:arg coefficient: UFL Coefficient
:arg num: coefficient index in the original form
:arg name: unique name to refer to the Coefficient in the kernel
:arg interior_facet: interior facet integral?
:returns: GEM expression referring to the Coefficient value
"""
varexp = gem.Variable(name, (None, None))
if coefficient.ufl_element().family() == 'Real':
size = numpy.prod(coefficient.ufl_shape, dtype=int)
data = gem.view(varexp, slice(num, num + 1), slice(size))
return gem.reshape(data, (), coefficient.ufl_shape)
element = create_element(coefficient.ufl_element())
size = numpy.prod(element.index_shape, dtype=int)
def expression(data):
result, = prune(
[gem.reshape(gem.view(data, slice(size)), element.index_shape)])
return result
if not interior_facet:
data = gem.view(varexp, slice(num, num + 1), slice(size))
return expression(gem.reshape(data, (), (size, )))
else:
data_p = gem.view(varexp, slice(num, num + 1), slice(size))
data_m = gem.view(varexp, slice(num, num + 1), slice(size, 2 * size))
return (expression(gem.reshape(data_p, (), (size, ))),
expression(gem.reshape(data_m, (), (size, ))))
def basis_evaluation(self, order, ps, entity=None):
"""Return code for evaluating the element at known points on the
reference element.
:param order: return derivatives up to this order.
:param ps: the point set object.
:param entity: the cell entity on which to tabulate.
"""
# Spatial dimension
dimension = self.cell.get_spatial_dimension()
# Shape of the tabulation matrix
shape = tuple(
index.extent
for index in ps.indices) + self.index_shape + self.value_shape
result = {}
for derivative in range(order + 1):
for alpha in mis(dimension, derivative):
name = str.format("rt_{}_{}_{}_{}_{}_{}", self.variant,
self.degree, ''.join(map(str, alpha)),
self.shift_axes,
'c' if self.continuous else 'd', {
None: "",
'+': "p",
'-': "m"
}[self.restriction])
result[alpha] = gem.partial_indexed(gem.Variable(name, shape),
ps.indices)
return result
def gem_to_loopy(gem_expr, var2terminal, scalar_type):
""" Method encapsulating stage 2.
Converts the gem expression dag into imperoc first, and then further into loopy.
:return slate_loopy: 2-tuple of loopy kernel for slate operations
and loopy GlobalArg for the output variable.
"""
# Creation of return variables for outer loopy
shape = gem_expr.shape if len(gem_expr.shape) != 0 else (1, )
idx = make_indices(len(shape))
indexed_gem_expr = gem.Indexed(gem_expr, idx)
args = ([
loopy.GlobalArg("output",
shape=shape,
dtype=scalar_type,
is_output=True,
is_input=True)
] + [
loopy.GlobalArg(var.name, shape=var.shape, dtype=scalar_type)
for var in var2terminal.keys()
])
ret_vars = [gem.Indexed(gem.Variable("output", shape), idx)]
preprocessed_gem_expr = impero_utils.preprocess_gem([indexed_gem_expr])
# glue assignments to return variable
assignments = list(zip(ret_vars, preprocessed_gem_expr))
# Part A: slate to impero_c
impero_c = impero_utils.compile_gem(assignments, (), remove_zeros=False)
# Part B: impero_c to loopy
return generate_loopy(impero_c, args, scalar_type, "slate_loopy",
[]), args[0].copy()
def prepare_arguments(arguments,
multiindices,
scalar_type,
interior_facet=False):
"""Bridges the kernel interface and the GEM abstraction for
Arguments. Vector Arguments are rearranged here for interior
facet integrals.
:arg arguments: UFL Arguments
:arg multiindices: Argument multiindices
:arg interior_facet: interior facet integral?
:returns: (funarg, expression)
funarg - :class:`coffee.Decl` function argument
expressions - GEM expressions referring to the argument
tensor
"""
assert isinstance(interior_facet, bool)
if len(arguments) == 0:
# No arguments
funarg = coffee.Decl(scalar_type, coffee.Symbol("A", rank=(1, )))
expression = gem.Indexed(gem.Variable("A", (1, )), (0, ))
return funarg, [expression]
elements = tuple(create_element(arg.ufl_element()) for arg in arguments)
shapes = tuple(element.index_shape for element in elements)
def expression(restricted):
return gem.Indexed(gem.reshape(restricted, *shapes),
tuple(chain(*multiindices)))
u_shape = numpy.array([numpy.prod(shape, dtype=int) for shape in shapes])
if interior_facet:
c_shape = tuple(2 * u_shape)
slicez = [[
slice(r * s, (r + 1) * s) for r, s in zip(restrictions, u_shape)
] for restrictions in product((0, 1), repeat=len(arguments))]
else:
c_shape = tuple(u_shape)
slicez = [[slice(s) for s in u_shape]]
funarg = coffee.Decl(scalar_type, coffee.Symbol("A", rank=c_shape))
varexp = gem.Variable("A", c_shape)
expressions = [expression(gem.view(varexp, *slices)) for slices in slicez]
return funarg, prune(expressions)
def __init__(self, integral_type, subdomain_id, domain_number):
"""Initialise a kernel builder."""
super(KernelBuilder,
self).__init__(integral_type.startswith("interior_facet"))
self.integral_type = integral_type
self.local_tensor = None
self.coordinates_args = None
self.coefficient_args = None
self.coefficient_split = None
if self.interior_facet:
self._cell_orientations = (gem.Variable("cell_orientation_0", ()),
gem.Variable("cell_orientation_1", ()))
else:
self._cell_orientations = (gem.Variable("cell_orientation", ()), )
if integral_type == "exterior_facet":
self._entity_number = {
None: gem.VariableIndex(gem.Variable("facet", ()))
}
elif integral_type == "interior_facet":
self._entity_number = {
'+': gem.VariableIndex(gem.Variable("facet_0", ())),
'-': gem.VariableIndex(gem.Variable("facet_1", ()))
}
elif integral_type == "vertex":
self._entity_number = {
None: gem.VariableIndex(gem.Variable("vertex", ()))
}
def __init__(self, scalar_type=None, interior_facet=False):
"""Initialise a kernel builder.
:arg interior_facet: kernel accesses two cells
"""
if scalar_type is None:
from tsfc.parameters import SCALAR_TYPE
scalar_type = SCALAR_TYPE
super(KernelBuilderBase, self).__init__(scalar_type=scalar_type,
interior_facet=interior_facet)
# Cell orientation
if self.interior_facet:
cell_orientations = gem.Variable("cell_orientations", (2, ))
self._cell_orientations = (gem.Indexed(cell_orientations, (0, )),
gem.Indexed(cell_orientations, (1, )))
else:
cell_orientations = gem.Variable("cell_orientations", (1, ))
self._cell_orientations = (gem.Indexed(cell_orientations, (0, )), )
def _coefficient(self, coefficient, name):
element = create_element(coefficient.ufl_element())
shape = self.shape + element.index_shape
size = numpy.prod(shape, dtype=int)
funarg = ast.Decl(SCALAR_TYPE, ast.Symbol(name), pointers=[("restrict", )],
qualifiers=["const"])
expression = gem.reshape(gem.Variable(name, (size, )), shape)
expression = gem.partial_indexed(expression, self.indices)
self.coefficient_map[coefficient] = expression
return funarg
def __init__(self, scalar_type, interior_facet=False):
"""Initialise a kernel builder.
:arg interior_facet: kernel accesses two cells
"""
super(KernelBuilderBase, self).__init__(scalar_type=scalar_type,
interior_facet=interior_facet)
# Cell orientation
if self.interior_facet:
shape = (2,)
cell_orientations = gem.Variable("cell_orientations", shape)
self._cell_orientations = (gem.Indexed(cell_orientations, (0,)),
gem.Indexed(cell_orientations, (1,)))
else:
shape = (1,)
cell_orientations = gem.Variable("cell_orientations", shape)
self._cell_orientations = (gem.Indexed(cell_orientations, (0,)),)
self.cell_orientations_loopy_arg = lp.GlobalArg("cell_orientations", dtype=numpy.int32, shape=shape)
def prepare_arguments(arguments,
multiindices,
scalar_type,
interior_facet=False):
"""Bridges the kernel interface and the GEM abstraction for
Arguments. Vector Arguments are rearranged here for interior
facet integrals.
:arg arguments: UFL Arguments
:arg multiindices: Argument multiindices
:arg interior_facet: interior facet integral?
:returns: (funarg, prepare, expressions)
funarg - :class:`coffee.Decl` function argument
prepare - list of COFFEE nodes to be prepended to the
kernel body
expressions - GEM expressions referring to the argument
tensor
"""
funarg = coffee.Decl(scalar_type, coffee.Symbol("A"), pointers=[()])
varexp = gem.Variable("A", (None, ))
if len(arguments) == 0:
# No arguments
zero = coffee.FlatBlock("memset({name}, 0, sizeof(*{name}));\n".format(
name=funarg.sym.gencode()))
return funarg, [zero], [gem.reshape(varexp, ())]
elements = tuple(create_element(arg.ufl_element()) for arg in arguments)
shapes = [element.index_shape for element in elements]
indices = tuple(chain(*multiindices))
def expression(restricted):
return gem.Indexed(gem.reshape(restricted, *shapes), indices)
u_shape = numpy.array(
[numpy.prod(element.index_shape, dtype=int) for element in elements])
if interior_facet:
c_shape = tuple(2 * u_shape)
slicez = [[
slice(r * s, (r + 1) * s) for r, s in zip(restrictions, u_shape)
] for restrictions in product((0, 1), repeat=len(arguments))]
else:
c_shape = tuple(u_shape)
slicez = [[slice(s) for s in u_shape]]
expressions = [
expression(gem.view(gem.reshape(varexp, c_shape), *slices))
for slices in slicez
]
zero = coffee.FlatBlock(
str.format("memset({name}, 0, {size} * sizeof(*{name}));\n",
name=funarg.sym.gencode(),
size=numpy.product(c_shape, dtype=int)))
return funarg, [zero], prune(expressions)
def test_cellwise_constant(cell, degree):
dim = cell.get_spatial_dimension()
element = finat.Lagrange(cell, degree)
index = gem.Index()
point = gem.partial_indexed(gem.Variable('X', (17, dim)), (index, ))
order = 2
for alpha, table in element.point_evaluation(order, point).items():
if sum(alpha) < degree:
assert table.free_indices == (index, )
else:
assert table.free_indices == ()
def prepare_coefficient(coefficient, name, scalar_type, interior_facet=False):
"""Bridges the kernel interface and the GEM abstraction for
Coefficients.
:arg coefficient: UFL Coefficient
:arg name: unique name to refer to the Coefficient in the kernel
:arg interior_facet: interior facet integral?
:returns: (funarg, expression)
funarg - :class:`loopy.GlobalArg` function argument
expression - GEM expression referring to the Coefficient
values
"""
assert isinstance(interior_facet, bool)
if coefficient.ufl_element().family() == 'Real':
# Constant
funarg = lp.GlobalArg(name,
dtype=scalar_type,
shape=(coefficient.ufl_element().value_size(), ))
expression = gem.reshape(gem.Variable(name, (None, )),
coefficient.ufl_shape)
return funarg, expression
finat_element = create_element(coefficient.ufl_element())
shape = finat_element.index_shape
size = numpy.prod(shape, dtype=int)
if not interior_facet:
expression = gem.reshape(gem.Variable(name, (size, )), shape)
else:
varexp = gem.Variable(name, (2 * size, ))
plus = gem.view(varexp, slice(size))
minus = gem.view(varexp, slice(size, 2 * size))
expression = (gem.reshape(plus, shape), gem.reshape(minus, shape))
size = size * 2
funarg = lp.GlobalArg(name, dtype=scalar_type, shape=(size, ))
return funarg, expression
def coefficient(self, o):
# Because we act on dofs, the ufl_shape is not the right thing to check
shape = o.dat.dim
try:
var = self.varmapping[o]
except KeyError:
name = f"C{len(self.varmapping)}"
var = gem.Variable(name, shape)
self.varmapping[o] = var
if o.ufl_shape == ():
assert shape == (1, )
return gem.Indexed(var, (0, ))
else:
return var
def gem_to_loopy(gem_expr):
""" Method encapsulating stage 2.
Converts the gem expression dag into imperoc first, and then further into loopy.
:return slate_loopy: loopy kernel for slate operations.
"""
# Creation of return variables for outer loopy
shape = gem_expr.shape if len(gem_expr.shape) != 0 else (1,)
idx = make_indices(len(shape))
indexed_gem_expr = gem.Indexed(gem_expr, idx)
arg = [loopy.GlobalArg("output", shape=shape)]
ret_vars = [gem.Indexed(gem.Variable("output", shape), idx)]
preprocessed_gem_expr = impero_utils.preprocess_gem([indexed_gem_expr])
# glue assignments to return variable
assignments = list(zip(ret_vars, preprocessed_gem_expr))
# Part A: slate to impero_c
impero_c = impero_utils.compile_gem(assignments, (), remove_zeros=False)
# Part B: impero_c to loopy
return generate_loopy(impero_c, arg, parameters["form_compiler"]["scalar_type"], "slate_loopy", [])
def rebuild_dg(element, expr, rt_var_name):
# To tabulate on the given element (which is on a different mesh to the
# expression) we must do so at runtime. We therefore create a quadrature
# element with runtime points to evaluate for each point in the element's
# dual basis. This exists on the same reference cell as the input element
# and we can interpolate onto it before mapping the result back onto the
# target space.
expr_tdim = expr.ufl_domain().topological_dimension()
# Need point evaluations and matching weights from dual basis.
# This could use FIAT's dual basis as below:
# num_points = sum(len(dual.get_point_dict()) for dual in element.fiat_equivalent.dual_basis())
# weights = []
# for dual in element.fiat_equivalent.dual_basis():
# pts = dual.get_point_dict().keys()
# for p in pts:
# for w, _ in dual.get_point_dict()[p]:
# weights.append(w)
# assert len(weights) == num_points
# but for now we just fix the values to what we know works:
if element.degree != 0 or not isinstance(element.cell,
FIAT.reference_element.Point):
raise NotImplementedError(
"Cross mesh interpolation only implemented for P0DG on vertex cells."
)
num_points = 1
weights = [1.] * num_points
# gem.Variable name starting with rt_ forces TSFC runtime tabulation
assert rt_var_name.startswith("rt_")
runtime_points_expr = gem.Variable(rt_var_name, (num_points, expr_tdim))
rule_pointset = finat.point_set.UnknownPointSet(runtime_points_expr)
try:
expr_fiat_cell = as_fiat_cell(expr.ufl_element().cell())
except AttributeError:
# expression must be pure function of spatial coordinates so
# domain has correct ufl cell
expr_fiat_cell = as_fiat_cell(expr.ufl_domain().ufl_cell())
rule = finat.quadrature.QuadratureRule(rule_pointset, weights=weights)
return finat.QuadratureElement(expr_fiat_cell, rule)
def __init__(self,
integral_type,
subdomain_id,
domain_number,
scalar_type=None,
diagonal=False):
"""Initialise a kernel builder."""
if diagonal:
raise NotImplementedError(
"Assembly of diagonal not implemented yet, sorry")
super(KernelBuilder,
self).__init__(scalar_type,
integral_type.startswith("interior_facet"))
self.integral_type = integral_type
self.local_tensor = None
self.coordinates_args = None
self.coefficient_args = None
self.coefficient_split = None
if self.interior_facet:
self._cell_orientations = (gem.Variable("cell_orientation_0", ()),
gem.Variable("cell_orientation_1", ()))
else:
self._cell_orientations = (gem.Variable("cell_orientation", ()), )
if integral_type == "exterior_facet":
self._entity_number = {
None: gem.VariableIndex(gem.Variable("facet", ()))
}
elif integral_type == "interior_facet":
self._entity_number = {
'+': gem.VariableIndex(gem.Variable("facet_0", ())),
'-': gem.VariableIndex(gem.Variable("facet_1", ()))
}
elif integral_type == "vertex":
self._entity_number = {
None: gem.VariableIndex(gem.Variable("vertex", ()))
}
def compile_element(expression,
dual_space=None,
parameters=None,
name="evaluate"):
"""Generate code for point evaluations.
:arg expression: A UFL expression (may contain up to one coefficient, or one argument)
:arg dual_space: if the expression has an argument, should we also distribute residual data?
:returns: Some coffee AST
"""
if parameters is None:
parameters = default_parameters()
else:
_ = default_parameters()
_.update(parameters)
parameters = _
expression = tsfc.ufl_utils.preprocess_expression(expression)
# # Collect required coefficients
try:
arg, = extract_coefficients(expression)
argument_multiindices = ()
coefficient = True
if expression.ufl_shape:
tensor_indices = tuple(gem.Index() for s in expression.ufl_shape)
else:
tensor_indices = ()
except ValueError:
arg, = extract_arguments(expression)
finat_elem = create_element(arg.ufl_element())
argument_multiindices = (finat_elem.get_indices(), )
argument_multiindex, = argument_multiindices
value_shape = finat_elem.value_shape
if value_shape:
tensor_indices = argument_multiindex[-len(value_shape):]
else:
tensor_indices = ()
coefficient = False
# Replace coordinates (if any)
builder = firedrake_interface.KernelBuilderBase(scalar_type=ScalarType_c)
domain = expression.ufl_domain()
# Translate to GEM
cell = domain.ufl_cell()
dim = cell.topological_dimension()
point = gem.Variable('X', (dim, ))
point_arg = ast.Decl(ScalarType_c, ast.Symbol('X', rank=(dim, )))
config = dict(interface=builder,
ufl_cell=cell,
precision=parameters["precision"],
point_indices=(),
point_expr=point,
argument_multiindices=argument_multiindices)
context = tsfc.fem.GemPointContext(**config)
# Abs-simplification
expression = tsfc.ufl_utils.simplify_abs(expression)
# Translate UFL -> GEM
if coefficient:
assert dual_space is None
f_arg = [builder._coefficient(arg, "f")]
else:
f_arg = []
translator = tsfc.fem.Translator(context)
result, = map_expr_dags(translator, [expression])
b_arg = []
if coefficient:
if expression.ufl_shape:
return_variable = gem.Indexed(
gem.Variable('R', expression.ufl_shape), tensor_indices)
result_arg = ast.Decl(ScalarType_c,
ast.Symbol('R', rank=expression.ufl_shape))
result = gem.Indexed(result, tensor_indices)
else:
return_variable = gem.Indexed(gem.Variable('R', (1, )), (0, ))
result_arg = ast.Decl(ScalarType_c, ast.Symbol('R', rank=(1, )))
else:
return_variable = gem.Indexed(
gem.Variable('R', finat_elem.index_shape), argument_multiindex)
result = gem.Indexed(result, tensor_indices)
if dual_space:
elem = create_element(dual_space.ufl_element())
if elem.value_shape:
var = gem.Indexed(gem.Variable("b", elem.value_shape),
tensor_indices)
b_arg = [
ast.Decl(ScalarType_c,
ast.Symbol("b", rank=elem.value_shape))
]
else:
var = gem.Indexed(gem.Variable("b", (1, )), (0, ))
b_arg = [ast.Decl(ScalarType_c, ast.Symbol("b", rank=(1, )))]
result = gem.Product(result, var)
result_arg = ast.Decl(ScalarType_c,
ast.Symbol('R', rank=finat_elem.index_shape))
# Unroll
max_extent = parameters["unroll_indexsum"]
if max_extent:
def predicate(index):
return index.extent <= max_extent
result, = gem.optimise.unroll_indexsum([result], predicate=predicate)
# Translate GEM -> COFFEE
result, = gem.impero_utils.preprocess_gem([result])
impero_c = gem.impero_utils.compile_gem([(return_variable, result)],
tensor_indices)
body = generate_coffee(impero_c, {}, parameters["precision"], ScalarType_c)
# Build kernel tuple
kernel_code = builder.construct_kernel(
"pyop2_kernel_" + name, [result_arg] + b_arg + f_arg + [point_arg],
body)
return kernel_code
def compile_expression_at_points(expression, points, coordinates, parameters=None):
"""Compiles a UFL expression to be evaluated at compile-time known
reference points. Useful for interpolating UFL expressions onto
function spaces with only point evaluation nodes.
:arg expression: UFL expression
:arg points: reference coordinates of the evaluation points
:arg coordinates: the coordinate function
:arg parameters: parameters object
"""
import coffee.base as ast
if parameters is None:
parameters = default_parameters()
else:
_ = default_parameters()
_.update(parameters)
parameters = _
# No arguments, please!
if extract_arguments(expression):
return ValueError("Cannot interpolate UFL expression with Arguments!")
# Apply UFL preprocessing
expression = ufl_utils.preprocess_expression(expression)
# Initialise kernel builder
builder = firedrake_interface.ExpressionKernelBuilder()
# Replace coordinates (if any)
domain = expression.ufl_domain()
if domain:
assert coordinates.ufl_domain() == domain
builder.domain_coordinate[domain] = coordinates
# Collect required coefficients
coefficients = extract_coefficients(expression)
if has_type(expression, GeometricQuantity):
coefficients = [coordinates] + coefficients
builder.set_coefficients(coefficients)
# Split mixed coefficients
expression = ufl_utils.split_coefficients(expression, builder.coefficient_split)
# Translate to GEM
point_set = PointSet(points)
config = dict(interface=builder,
ufl_cell=coordinates.ufl_domain().ufl_cell(),
precision=parameters["precision"],
point_set=point_set)
ir, = fem.compile_ufl(expression, point_sum=False, **config)
# Deal with non-scalar expressions
value_shape = ir.shape
tensor_indices = tuple(gem.Index() for s in value_shape)
if value_shape:
ir = gem.Indexed(ir, tensor_indices)
# Build kernel body
return_shape = (len(points),) + value_shape
return_indices = point_set.indices + tensor_indices
return_var = gem.Variable('A', return_shape)
return_arg = ast.Decl(SCALAR_TYPE, ast.Symbol('A', rank=return_shape))
return_expr = gem.Indexed(return_var, return_indices)
ir, = impero_utils.preprocess_gem([ir])
impero_c = impero_utils.compile_gem([(return_expr, ir)], return_indices)
point_index, = point_set.indices
body = generate_coffee(impero_c, {point_index: 'p'}, parameters["precision"])
# Handle cell orientations
if builder.needs_cell_orientations([ir]):
builder.require_cell_orientations()
# Build kernel tuple
return builder.construct_kernel(return_arg, body)
def prepare_coefficients(coefficients, num, name, interior_facet=False):
"""Bridges the kernel interface and the GEM abstraction for
Coefficients.
:arg coefficients: split UFL Coefficients
:arg num: coefficient index in the original form
:arg name: unique name to refer to the Coefficient in the kernel
:arg interior_facet: interior facet integral?
:returns: GEM expression referring to the Coefficient value
"""
varexp = gem.Variable(name, (None, None))
if len(coefficients) == 1 and coefficients[0].ufl_element().family() == 'Real':
coefficient, = coefficients
size = numpy.prod(coefficient.ufl_shape, dtype=int)
data = gem.view(varexp, slice(num, num + 1), slice(size))
expression = gem.reshape(data, (), coefficient.ufl_shape)
return [expression]
elements = [create_element(coeff.ufl_element()) for coeff in coefficients]
space_dimensions = [numpy.prod(element.index_shape, dtype=int)
for element in elements]
ends = list(numpy.cumsum(space_dimensions))
starts = [0] + ends[:-1]
slices = [slice(start, end) for start, end in zip(starts, ends)]
transposed_shapes = []
tensor_ranks = []
for element in elements:
if isinstance(element, TensorFiniteElement):
scalar_shape = element.base_element.index_shape
tensor_shape = element.index_shape[len(scalar_shape):]
else:
scalar_shape = element.index_shape
tensor_shape = ()
transposed_shapes.append(tensor_shape + scalar_shape)
tensor_ranks.append(len(tensor_shape))
def transpose(expr, rank):
assert not expr.free_indices
assert 0 <= rank < len(expr.shape)
if rank == 0:
return expr
else:
indices = tuple(gem.Index(extent=extent) for extent in expr.shape)
transposed_indices = indices[rank:] + indices[:rank]
return gem.ComponentTensor(gem.Indexed(expr, indices),
transposed_indices)
def expressions(data):
return prune([transpose(gem.reshape(gem.view(data, slice_), shape), rank)
for slice_, shape, rank in zip(slices, transposed_shapes, tensor_ranks)])
size = sum(space_dimensions)
if not interior_facet:
data = gem.view(varexp, slice(num, num + 1), slice(size))
return expressions(gem.reshape(data, (), (size,)))
else:
data_p = gem.view(varexp, slice(num, num + 1), slice(size))
data_m = gem.view(varexp, slice(num, num + 1), slice(size, 2 * size))
return list(zip(expressions(gem.reshape(data_p, (), (size,))),
expressions(gem.reshape(data_m, (), (size,)))))
def vector():
return gem.Variable('u', (12, ))
def matrix():
return gem.Variable('A', (10, 12))
