-
Marcel Koch authoredMarcel Koch authored
from dune.codegen.tools import get_pymbolic_basename
from dune.codegen.generation import (iname,
domain,
temporary_variable,
instruction, globalarg, preamble)
from dune.codegen.pdelab.geometry import world_dimension
from dune.codegen.options import get_form_option
import pymbolic.primitives as prim
# add inames over the micro elements in tensor representation,
# i.e. each element has (i_1,i_2,...,i_d) indices
@iname
def sub_element_inames():
name = "subel"
dim = world_dimension()
dim_names = ["x", "y", "z"] + [str(i) for i in range(4, dim + 1)]
inames = tuple()
for i in range(dim):
inames = inames + ("subel_" + dim_names[i],)
domain("subel_" + dim_names[i], get_form_option("number_of_blocks"))
return inames
# compute sequential index for given tensor index, the same as index in base-k to base-10
def tensor_index_to_sequential_index(indices, k):
return prim.Sum(tuple(prim.Variable(index) * k ** i for i, index in enumerate(indices)))
# compute tensor index for given sequential index, the same as index in base-10 to base-k
def sequential_index_to_tensor_index(iname, k):
return tuple(prim.Remainder(prim.Variable(iname) / (k**i), k) for i in range(world_dimension()))
# compute index for higher order FEM for a given Q1 index
def micro_index_to_macro_index(element, inames):
subelem_inames = sub_element_inames()
k = get_form_option("number_of_blocks")
p = element.degree()
return prim.Sum(tuple((p * prim.Variable(si) + prim.Variable(bi)) * (p * k + 1) ** i
for i, (si, bi) in enumerate(zip(subelem_inames, inames))))
# translate a point in the micro element into macro coordinates
def define_point_in_macro(name, point_in_micro, visitor):
# TODO this won't work for 2d mannifolds
dim = world_dimension()
if get_form_option('vectorization_blockstructured'):
temporary_variable(name, shape=(dim,), managed=True)
else:
temporary_variable(name, shape=(dim,), shape_impl=('fv',))
# point_macro = (point_micro + index_micro) / number_of_blocks
# where index_micro = tensor index of the micro element
subelem_inames = sub_element_inames()
for i in range(dim):
if isinstance(point_in_micro, prim.Subscript):
expr = prim.Subscript(point_in_micro.aggregate, point_in_micro.index + (i,))
else:
expr = prim.Subscript(point_in_micro, (i,))
expr = prim.Sum((expr, prim.Variable(subelem_inames[i]),))
expr = prim.Quotient(expr, get_form_option('number_of_blocks'))
# TODO relax within inames
instruction(assignee=prim.Subscript(prim.Variable(name), (i,)),
expression=expr,
within_inames=frozenset(subelem_inames + visitor.quadrature_inames()),
tags=frozenset({subelem_inames[i]})
)
def name_point_in_macro(point_in_micro, visitor):
assert isinstance(point_in_micro, prim.Expression)
name = get_pymbolic_basename(point_in_micro) + "_macro"
define_point_in_macro(name, point_in_micro, visitor)
return name
@preamble
def define_container_alias(name, container, lfs, element, is_const):
k = get_form_option("number_of_blocks")
p = element.degree()
dim = world_dimension()
element_stride = tuple(p * (p * k + 1)**i for i in range(0, dim))
index_stride = tuple((p * k + 1)**i for i in range(0, dim))
globalarg(name, shape=(k,) * dim + (p + 1,) * dim, strides=element_stride + index_stride, managed=True)
if is_const:
return "const auto {} = &{}({},0);".format(name, container, lfs.name)
else:
return "auto {} = &{}.container()({},0);".format(name, container, lfs.name)
def name_accumulation_alias(container, accumspace):
name = container + "_" + accumspace.lfs.name + "_alias"
name_tail = container + "_" + accumspace.lfs.name + "_alias_tail"
define_container_alias(name, container, accumspace.lfs, accumspace.element, is_const=False)
define_container_alias(name_tail, container, accumspace.lfs, accumspace.element, is_const=False)
return name
def name_container_alias(container, lfs, element):
name = container + "_" + lfs.name + "_alias"
define_container_alias(name, container, lfs, element, is_const=True)
return name