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_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 prepare_arguments(arguments, multiindices, 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 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 prepare_arguments(arguments, multiindices, 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 expression(data):
result, = prune(
[gem.reshape(gem.view(data, slice(size)), element.index_shape)])
return result
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)])
def expression(data):
result, = prune([gem.reshape(gem.view(data, slice(size)), element.index_shape)])
return result
