diff --git a/python/dune/perftool/loopy/transformer.py b/python/dune/perftool/loopy/transformer.py
index 0a0fde9abc31dbd7447a53e1918084b16df5c9de..9dec35bc7bb3e39d438a0d441a0f25e82b5be471 100644
--- a/python/dune/perftool/loopy/transformer.py
+++ b/python/dune/perftool/loopy/transformer.py
@@ -31,11 +31,11 @@ from pymbolic.primitives import Subscript, Variable
class UFL2LoopyVisitor(ModifiedTerminalTracker, UFL2PymbolicMapper, GeometryMapper):
- def __init__(self, measure, subdomain_id, dimension_index_aliases):
+ def __init__(self, measure, subdomain_id, dimension_indices):
# Some variables describing the integral measure of this integral
self.measure = measure
self.subdomain_id = subdomain_id
- self.dimension_index_aliases = dimension_index_aliases
+ self.dimension_indices = dimension_indices
# Call base class constructors
super(UFL2LoopyVisitor, self).__init__()
@@ -96,7 +96,8 @@ class UFL2LoopyVisitor(ModifiedTerminalTracker, UFL2PymbolicMapper, GeometryMapp
if len(subel.value_shape()) != 0:
from dune.perftool.pdelab.geometry import dimension_iname
from dune.perftool.pdelab.basis import lfs_child
- lfs = lfs_child(lfs, dimension_iname(context='arg'))
+ idims = tuple(dimension_iname(context='arg', count=i) for i in range(len(subel.value_shape())))
+ lfs = lfs_child(lfs, idims, shape=subel.value_shape())
residual_shape[icount] = name_lfs_bound(name_leaf_lfs(subel.sub_elements()[0], ma.restriction))
subel = subel.sub_elements()[0]
else:
@@ -181,8 +182,8 @@ class UFL2LoopyVisitor(ModifiedTerminalTracker, UFL2PymbolicMapper, GeometryMapp
# Now treat the case of this being a vector finite element
if element.num_sub_elements() > 0:
# I cannot handle general mixed elements here...
- from ufl import VectorElement
- assert isinstance(element, VectorElement)
+ from ufl import VectorElement, TensorElement
+ assert isinstance(element, (VectorElement, TensorElement))
# Determine whether this is a non-scalar subargument. This information is later
# used by the Indexed-Mapper to make some indices loop indices!
@@ -191,7 +192,7 @@ class UFL2LoopyVisitor(ModifiedTerminalTracker, UFL2PymbolicMapper, GeometryMapp
# If this is a vector element, we need add an additional accumulation loop iname
from dune.perftool.pdelab.geometry import dimension_iname
for i in range(self.argshape):
- self.inames.append(dimension_iname(context='arg'))
+ self.inames.append(dimension_iname(context='arg', count=i))
# For the purpose of basis evaluation, we need to take the leaf element
leaf_element = element.sub_elements()[0]
@@ -286,7 +287,7 @@ class UFL2LoopyVisitor(ModifiedTerminalTracker, UFL2PymbolicMapper, GeometryMapp
# Recurse to get the summation expression
term = self.call(o.ufl_operands[0])
- self.redinames = tuple(i for i in self.redinames if i not in self.dimension_index_aliases)
+ self.redinames = tuple(i for i in self.redinames if i not in self.dimension_indices)
if len(self.redinames) > 0:
ret = Reduction("sum", tuple(name_index(ind) for ind in self.redinames), term)
name = get_temporary_name()
@@ -314,10 +315,10 @@ class UFL2LoopyVisitor(ModifiedTerminalTracker, UFL2PymbolicMapper, GeometryMapp
else:
from pymbolic.primitives import Variable
from dune.perftool.pdelab import name_index
- if index in self.dimension_index_aliases:
+ if index in self.dimension_indices:
from dune.perftool.pdelab.geometry import dimension_iname
- self.inames.append(dimension_iname(context='arg'))
- return Variable(dimension_iname(context='arg'))
+ self.inames.append(self.dimension_indices[index])
+ return Variable(self.dimension_indices[index])
else:
return Variable(name_index(index))
diff --git a/python/dune/perftool/pdelab/basis.py b/python/dune/perftool/pdelab/basis.py
index 9a45e0d297efe75c63e5e24c451a183b47d5cb41..7e001b4af62a7c1c650158ce19c089595e5d075a 100644
--- a/python/dune/perftool/pdelab/basis.py
+++ b/python/dune/perftool/pdelab/basis.py
@@ -75,11 +75,26 @@ def define_lfs(name, father, child):
return "auto {} = child({}, _{});".format(name, father, child)
-def lfs_child(lfs, child):
- from pymbolic.primitives import Call
+def lfs_child(lfs, children, shape=None):
+ from pymbolic.primitives import Call, Product, Sum
+ # Old pre-TensorElement implementation kept for comaptibility
+ if shape is None:
+ indices = (Variable(children[0]),)
+ else:
+ children = tuple(Variable(c) for c in children)
+ offset = Product(tuple())
+ indices = (Sum(
+ tuple(
+ Product(
+ (Product(tuple(shape[j] for j in range(i))
+ ),
+ child)
+ )
+ for i, child in enumerate(children))
+ ),)
+
from dune.perftool.loopy.functions import LFSChild
- return Call(LFSChild(lfs), (Variable(child),))
-# return "{}.child({})".format(lfs, child)
+ return Call(LFSChild(lfs), indices)
@generator_factory(cache_key_generator=lambda e, r, **kw: (e, r))
@@ -307,6 +322,14 @@ def name_basis_gradient(leaf_element, restriction):
return name
+def shape_as_pymbolic(shape):
+ def _shape_as_pymbolic(s):
+ if isinstance(s, str):
+ return Variable(s)
+ else:
+ return s
+ return tuple(_shape_as_pymbolic(s) for s in shape)
+
@cached
def evaluate_trialfunction(element, name, restriction, component):
from ufl.functionview import select_subelement
@@ -314,28 +337,28 @@ def evaluate_trialfunction(element, name, restriction, component):
# Determine the rank of the trialfunction tensor
rank = len(sub_element.value_shape())
- assert rank in (0, 1)
- shape = (name_dimension(),) * rank
- shape_impl = ('fv', ) * rank
+ shape = sub_element.value_shape()
+ shape_impl = ('arr', ) * rank
idims = tuple(dimension_iname(count=i) for i in range(rank))
leaf_element = sub_element
- from ufl import VectorElement
- if isinstance(sub_element, VectorElement):
+ from ufl import VectorElement, TensorElement
+ if isinstance(sub_element, (VectorElement, TensorElement)):
leaf_element = sub_element.sub_elements()[0]
temporary_variable(name, shape=shape, shape_impl=shape_impl)
lfs = name_lfs(element, restriction, component)
index = lfs_iname(leaf_element, restriction, context='trial')
basis = name_basis(leaf_element, restriction)
- assignee = Variable(name)
- if isinstance(sub_element, VectorElement):
- lfs = lfs_child(lfs, idims[0])
- assignee = Subscript(assignee, Variable(idims[0]))
+ if isinstance(sub_element, (VectorElement, TensorElement)):
+ # TODO we need a concept of symmetry here
+ lfs = lfs_child(lfs, idims, shape=shape_as_pymbolic(shape))
from dune.perftool.pdelab.argument import pymbolic_coefficient
coeff = pymbolic_coefficient(lfs, index, restriction)
+ assignee = Subscript(Variable(name), tuple(Variable(i) for i in idims))
+
reduction_expr = Product((coeff, Subscript(Variable(basis), Variable(index))))
instruction(expression=Reduction("sum", index, reduction_expr, allow_simultaneous=True),
assignee=assignee,
@@ -350,18 +373,15 @@ def evaluate_trialfunction_gradient(element, name, restriction, component):
from ufl.functionview import select_subelement
sub_element = select_subelement(element, component)
rank = len(sub_element.value_shape()) + 1
- assert rank in (1, 2)
# We do then set some variables accordingly
- shape = (name_dimension(),) * rank
- if rank == 1:
- shape_impl = ('fv',)
- else:
- shape_impl = ('fm',)
+ shape = sub_element.value_shape() + (element.cell().geometric_dimension(),)
+ shape_impl = ('arr',) * rank
+
idims = tuple(dimension_iname(count=i) for i in range(rank))
leaf_element = sub_element
- from ufl import VectorElement
- if isinstance(sub_element, VectorElement):
+ from ufl import VectorElement, TensorElement
+ if isinstance(sub_element, (VectorElement, TensorElement)):
leaf_element = sub_element.sub_elements()[0]
# and proceed to call the necessary generator functions
@@ -370,8 +390,9 @@ def evaluate_trialfunction_gradient(element, name, restriction, component):
index = lfs_iname(leaf_element, restriction, context='trialgrad')
basis = name_basis_gradient(leaf_element, restriction)
- if isinstance(sub_element, VectorElement):
- lfs = lfs_child(lfs, idims[0])
+ if isinstance(sub_element, (VectorElement, TensorElement)):
+ # TODO we need a concept of symmetry here
+ lfs = lfs_child(lfs, idims[:-1], shape=shape_as_pymbolic(shape[:-1]))
from dune.perftool.pdelab.argument import pymbolic_coefficient
coeff = pymbolic_coefficient(lfs, index, restriction)
diff --git a/python/dune/perftool/pdelab/driver.py b/python/dune/perftool/pdelab/driver.py
index 33d179fd7fc2b06924fa5e0470c65a96e0410007..a32841f879f0dcba39cf229ae5a2a94c22536a6a 100644
--- a/python/dune/perftool/pdelab/driver.py
+++ b/python/dune/perftool/pdelab/driver.py
@@ -73,11 +73,17 @@ def isQk(fem):
def isDG(fem):
return fem._short_name is 'DG'
+def flatten_tensor(expr):
+ from ufl import TensorElement
+ assert isinstance(expr, TensorElement)
+
def FEM_name_mangling(fem):
- from ufl import MixedElement, VectorElement, FiniteElement
+ from ufl import MixedElement, VectorElement, FiniteElement, TensorElement
if isinstance(fem, VectorElement):
return FEM_name_mangling(fem.sub_elements()[0]) + "_" + str(fem.num_sub_elements())
+ if isinstance(fem, TensorElement):
+ return FEM_name_mangling(fem.sub_elements()[0]) + "_" + "_".join(str(i) for i in fem.value_shape())
if isinstance(fem, MixedElement):
name = ""
for elem in fem.sub_elements():
@@ -413,8 +419,8 @@ def define_compositegfs_constraints(name, *children):
@symbol
def name_bctype_function(expr):
- from ufl import MixedElement, VectorElement
- if isinstance(expr, VectorElement):
+ from ufl import MixedElement, VectorElement, TensorElement
+ if isinstance(expr, (VectorElement, TensorElement)):
element, (dirichlet, _) = get_constraints_metadata(expr)
# get the correct name from the corresponding UFL file
from dune.perftool.generation import get_global_context_value
@@ -461,16 +467,17 @@ def name_assembled_constraints(expr):
@preamble
def typedef_gfs(element, dirichlet, name):
vb = type_vectorbackend()
- from ufl import FiniteElement, MixedElement, VectorElement, EnrichedElement, RestrictedElement
+ from ufl import FiniteElement, MixedElement, VectorElement, EnrichedElement, RestrictedElement, TensorElement
if isinstance(element, FiniteElement):
gv = type_leafview()
fem = type_fem(element)
cass = type_constraintsassembler(dirichlet)
return "typedef Dune::PDELab::GridFunctionSpace<{}, {}, {}, {}> {};".format(gv, fem, cass, vb, name)
- if isinstance(element, VectorElement):
+ if isinstance(element, (VectorElement, TensorElement)):
include_file("dune/pdelab/gridfunctionspace/vectorgridfunctionspace.hh", filetag="driver")
gv = type_leafview()
fem = type_fem(element.sub_elements()[0])
+ # TODO implement symmetry here
dim = element.num_sub_elements()
cass = type_constraintsassembler(dirichlet)
return "typedef Dune::PDELab::VectorGridFunctionSpace<{}, {}, {}, {}, {}, {}> {};".format(gv, fem, dim, vb, vb, cass, name)
@@ -495,12 +502,12 @@ def type_gfs(expr):
@preamble
def define_gfs(element, dirichlet, name):
gfstype = type_gfs(element)
- from ufl import FiniteElement, MixedElement, VectorElement, EnrichedElement, RestrictedElement
+ from ufl import FiniteElement, MixedElement, VectorElement, EnrichedElement, RestrictedElement, TensorElement
if isinstance(element, FiniteElement):
gv = name_leafview()
fem = name_fem(element)
return "{} {}({}, {});".format(gfstype, name, gv, fem)
- if isinstance(element, VectorElement):
+ if isinstance(element, (VectorElement, TensorElement)):
gv = name_leafview()
fem = name_fem(element.sub_elements()[0])
return "{} {}({}, {});".format(gfstype, name, gv, fem)
@@ -705,8 +712,8 @@ def define_compositegfs_parameterfunction(name, *children):
@symbol
def name_boundary_function(expr):
- from ufl import MixedElement, VectorElement
- if isinstance(expr, VectorElement):
+ from ufl import MixedElement, VectorElement, TensorElement
+ if isinstance(expr, (VectorElement, TensorElement)):
element, (_, boundary) = get_constraints_metadata(expr)
from dune.perftool.generation import get_global_context_value
name = get_global_context_value("namedata").get(id(boundary), "zero")
@@ -881,8 +888,8 @@ def print_matrix():
@preamble
def define_gfs_name(element, prefix=()):
- from ufl import MixedElement, VectorElement
- if isinstance(element, VectorElement):
+ from ufl import MixedElement, VectorElement, TensorElement
+ if isinstance(element, (VectorElement, TensorElement)):
gfs = name_gfs(element)
return "{}.name(\"data_{}\");".format(gfs, "_".join(str(p) for p in prefix))
if isinstance(element, MixedElement):
diff --git a/python/dune/perftool/pdelab/localoperator.py b/python/dune/perftool/pdelab/localoperator.py
index 42a9d4e124dc6506b4d17f72bd5cded9925352c5..85fef5fa2781caaceb50f9dde0dd6c3b6bdf26bf 100644
--- a/python/dune/perftool/pdelab/localoperator.py
+++ b/python/dune/perftool/pdelab/localoperator.py
@@ -161,8 +161,8 @@ def generate_kernel(integrals):
subdomain_id = integral.subdomain_id()
subdomain_data = integral.subdomain_data()
- from dune.perftool.ufl.dimensionindex import collect_dimension_index_aliases
- dimension_index_aliases = collect_dimension_index_aliases(integrand)
+ from dune.perftool.ufl.dimensionindex import dimension_index_mapping
+ dimension_indices = dimension_index_mapping(integrand)
# Now split the given integrand into accumulation expressions
from dune.perftool.ufl.transformations.extract_accumulation_terms import split_into_accumulation_terms
@@ -170,7 +170,7 @@ def generate_kernel(integrals):
# Get a transformer instance for this kernel
from dune.perftool.loopy.transformer import UFL2LoopyVisitor
- visitor = UFL2LoopyVisitor(measure, subdomain_id, dimension_index_aliases)
+ visitor = UFL2LoopyVisitor(measure, subdomain_id, dimension_indices)
# Iterate over the terms and generate a kernel
for term in accterms:
diff --git a/python/dune/perftool/ufl/dimensionindex.py b/python/dune/perftool/ufl/dimensionindex.py
index 1446bd4dc8f9b71fcc27d884b3af513c634eaf66..e7c48a25581fb950095ce121514336ef14212944 100644
--- a/python/dune/perftool/ufl/dimensionindex.py
+++ b/python/dune/perftool/ufl/dimensionindex.py
@@ -3,30 +3,32 @@
from ufl.algorithms import MultiFunction
-class _CollectDimensionIndexAliases(MultiFunction):
+class _DimensionIndexMapping(MultiFunction):
call = MultiFunction.__call__
def __call__(self, o):
self.shape = 0
- return self.call(o)
+ self.dimension_index_dict = {}
+ self.call(o)
- def expr(self, o):
- return frozenset({}).union(*tuple(self.call(op) for op in o.ufl_operands))
+ return self.dimension_index_dict
- def terminal(self, o):
- return frozenset({})
+ def expr(self, o):
+ for op in o.ufl_operands:
+ self.call(op)
def function_view(self, o):
- self.shape = len(o.ufl_operands[1])
- return frozenset({})
+ from ufl.functionview import select_subelement
+ subelement = select_subelement(o.ufl_operands[0].ufl_element(), o.ufl_operands[1])
+ self.shape = len(subelement.value_shape())
def indexed(self, o):
- ret = self.call(o.ufl_operands[0])
- if self.shape:
- ret = ret.union(frozenset({o.ufl_operands[1][:self.shape][0]}))
+ self.call(o.ufl_operands[0])
+ for i in range(self.shape):
+ from dune.perftool.pdelab.geometry import dimension_iname
+ self.dimension_index_dict[o.ufl_operands[1][i]] = dimension_iname(context='arg', count=i)
self.shape = 0
- return ret
-def collect_dimension_index_aliases(expr):
- return _CollectDimensionIndexAliases()(expr)
+def dimension_index_mapping(expr):
+ return _DimensionIndexMapping()(expr)