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