WIP: Changes the single loop over basis function into nested loops

python/dune/perftool/blockstructured/__init__.py

@@ -11,7 +11,8 @@ from dune.perftool.blockstructured.geometry import (pymbolic_jacobian_inverse_tr
                                                     pymbolic_jacobian_determinant,
                                                     pymbolic_facet_jacobian_determinant,
                                                     to_global)
-from dune.perftool.blockstructured.basis import pymbolic_reference_gradient
+from dune.perftool.blockstructured.basis import (pymbolic_reference_gradient,
+                                                 pymbolic_basis)
 from dune.perftool.blockstructured.tools import sub_element_inames
 
 from dune.perftool.pdelab import PDELabInterface
@@ -29,6 +30,17 @@ class BlockStructuredInterface(PDELabInterface):
     def lfs_inames(self, element, restriction, number=None, context=''):
         return lfs_inames(element, restriction, number, context) + sub_element_inames()
 
+
+    #
+    # Test and trial function related generator functions
+    #
+
+    def pymbolic_basis(self, element, restriction, number):
+        return pymbolic_basis(element, restriction, number)
+
+    def pymbolic_reference_gradient(self, element, restriction, number):
+        return pymbolic_reference_gradient(element, restriction, number)
+
     #
     # Geometry related generator functions
     #

python/dune/perftool/blockstructured/basis.py

@@ -8,16 +8,22 @@ from dune.perftool.generation import (backend,
                                       initializer_list,
                                       include_file,)
 from dune.perftool.tools import get_pymbolic_basename
+# TODO clean up some imports
 from dune.perftool.pdelab.basis import (declare_cache_temporary,
                                         name_localbasis_cache,
-                                        type_localbasis
-                                        )
+                                        type_localbasis,
+                                        restricted_name,
+                                        FEM_name_mangling)
 from dune.perftool.pdelab.driver import (basetype_range,
                                          isPk,
                                          isQk)
 from dune.perftool.pdelab.geometry import world_dimension
 from dune.perftool.pdelab.quadrature import pymbolic_quadrature_position_in_cell
 from dune.perftool.pdelab.spaces import type_gfs
+from dune.perftool.blockstructured.spaces import lfs_inames
+from dune.perftool.blockstructured.tools import tensor_index_to_sequential_index
+import pymbolic.primitives as prim
+
 
 
 @backend(interface="typedef_localbasis", name="blockstructured")
@@ -46,6 +52,19 @@ def typedef_localbasis(element, name):
     return "using {} = Dune::{};".format(name, basis_type)
 
 
+@class_member(classtag="operator")
+def define_localbasis(leaf_element, name):
+    localBasis_type = type_localbasis(leaf_element)
+    initializer_list(name, (), classtag="operator")
+    return "const {} {};".format(localBasis_type, name)
+
+
+def name_localbasis(leaf_element):
+    name = "microElementBasis"
+    globalarg(name)
+    define_localbasis(leaf_element, name)
+    return name
+
 
 @backend(interface="evaluate_basis", name="blockstructured")
 @kernel_cached
@@ -61,6 +80,17 @@ def evaluate_basis(leaf_element, name, restriction):
                 )
 
 
+@backend(interface="pymbolic_basis", name="blockstructured")
+def pymbolic_basis(leaf_element, restriction, number, context=''):
+    assert leaf_element.num_sub_elements() == 0
+    name = "phi_{}".format(FEM_name_mangling(leaf_element))
+    name = restricted_name(name, restriction)
+    evaluate_basis(leaf_element, name, restriction)
+    iname = lfs_inames(leaf_element, restriction, number, context=context)
+
+    return prim.Subscript(prim.Variable(name), (tensor_index_to_sequential_index(iname, 2), ))
+
+
 @backend(interface="evaluate_grad", name="blockstructured")
 @kernel_cached
 def evaluate_reference_gradient(leaf_element, name, restriction):
@@ -75,15 +105,12 @@ def evaluate_reference_gradient(leaf_element, name, restriction):
                 )
 
 
-@class_member(classtag="operator")
-def define_localbasis(leaf_element, name):
-    localBasis_type = type_localbasis(leaf_element)
-    initializer_list(name, (), classtag="operator")
-    return "const {} {};".format(localBasis_type, name)
-
+@backend(interface="pymbolic_reference_gradient", name="blockstructured")
+def pymbolic_reference_gradient(leaf_element, restriction, number, context=''):
+    assert leaf_element.num_sub_elements() == 0
+    name = "js_{}".format(FEM_name_mangling(leaf_element))
+    name = restricted_name(name, restriction)
+    evaluate_reference_gradient(leaf_element, name, restriction)
+    inames = lfs_inames(leaf_element, restriction, number, context=context)
 
-def name_localbasis(leaf_element):
-    name = "microElementBasis"
+    return prim.Subscript(prim.Variable(name), (tensor_index_to_sequential_index(inames,2), 0))
\ No newline at end of file
-    globalarg(name)
-    define_localbasis(leaf_element, name)
-    return name

python/dune/perftool/blockstructured/spaces.py

 from dune.perftool.generation import (backend,
                                       domain)
 from dune.perftool.pdelab.geometry import world_dimension
-from dune.perftool.blockstructured.tools import sub_element_inames
-
 
 #TODO subelem_inames hinzufugen?
 @backend(interface="lfs_inames", name="blockstructured")
@@ -14,7 +12,11 @@ def lfs_inames(element, restriction, count=None, context=''):
         else:
             context = str(count)
 
-    name = "micro_{}_index".format(context)
-    domain(name, pow(2,world_dimension()))
+    inames = tuple()
+    dim = world_dimension()
+    for i in range(dim):
+        name = "micro_{}_index_{}".format(context, i)
+        inames = inames + (name,)
+        domain(name, 2)
 
-    return (name, ) + sub_element_inames()
+    return inames

python/dune/perftool/blockstructured/tools.py

@@ -27,59 +27,63 @@ def sub_element_inames():
     return inames
 
 
+def tensor_index_to_sequential_index(inames, k):
+    index = prim.Sum(tuple(prim.Variable(name)*k**i for i,name in enumerate(inames)))
+    return index
+
+
 # TODO 3d
-def micro_index_to_macro_index(index):
-    if not isinstance(index, prim.Variable):
-        index = prim.Variable(index)
+def micro_index_to_macro_index(inames):
+    inames = tensor_index_to_sequential_index(inames, get_option("number_of_blocks")+1)
 
-    subelem_inames = sub_element_inames()
+    subelem_inames_orig = sub_element_inames()
 
     k = get_option("number_of_blocks")
 
     it = get_global_context_value("integral_type")
     if it == "cell":
-        x = subelem_inames[0]
-        if local_dimension() > 1:
-            y = subelem_inames[1]
-            if world_dimension() > 2:
-                z = subelem_inames[2]
+        sublem_inames = subelem_inames_orig
     elif it == "exterior_facet" or it == "interior_facet":
         center = pymbolic_in_cell_coordinates(prim.Variable(name_localcenter()), Restriction.NEGATIVE)
 
         def predicate(x):
             return frozenset([prim.Comparison(prim.Subscript(center, 0), x, 0.5)])
 
-        x = "x"
-        temporary_variable(x)
-        instruction(assignee=prim.Variable(x),
-                    expression=prim.Variable(subelem_inames[0]),
-                    within_inames=frozenset(subelem_inames).union(frozenset(quadrature_inames())),
+        sublem_inames = ("x",)
+        temporary_variable(sublem_inames[0])
+        instruction(assignee=prim.Variable(sublem_inames[0]),
+                    expression=prim.Variable(subelem_inames_orig[0]),
+                    within_inames=frozenset(subelem_inames_orig).union(frozenset(quadrature_inames())),
                     predicates=predicate("==")
                     )
-        instruction(assignee=prim.Variable(x),
+        instruction(assignee=prim.Variable(sublem_inames[0]),
                     expression=prim.Product(((k-1),prim.Subscript(center, (0,)))),
-                    within_inames=frozenset(subelem_inames).union(frozenset(quadrature_inames())),
+                    within_inames=frozenset(subelem_inames_orig).union(frozenset(quadrature_inames())),
                     predicates=predicate("!=")
                     )
-        y = "y"
-        temporary_variable(y)
-        instruction(assignee=prim.Variable(y),
+        sublem_inames = sublem_inames + ("y",)
+        temporary_variable(sublem_inames[1])
+        instruction(assignee=prim.Variable(sublem_inames[1]),
                     expression=prim.Product(((k-1),prim.Subscript(center, (1,)))),
-                    within_inames=frozenset(subelem_inames).union(frozenset(quadrature_inames())),
+                    within_inames=frozenset(subelem_inames_orig).union(frozenset(quadrature_inames())),
                     predicates=predicate("==")
                     )
-        instruction(assignee=prim.Variable(y),
-                    expression=prim.Variable(subelem_inames[0]),
-                    within_inames=frozenset(subelem_inames).union(frozenset(quadrature_inames())),
+        instruction(assignee=prim.Variable(sublem_inames[1]),
+                    expression=prim.Variable(subelem_inames_orig[0]),
+                    within_inames=frozenset(subelem_inames_orig).union(frozenset(quadrature_inames())),
                     predicates=predicate("!=")
                     )
+    #
+    # modified_index = prim.Sum((prim.Variable(x), index))
+    # if local_dimension() > 1:
+    #     modified_index = prim.Sum((modified_index, prim.Product(((k+1), prim.Variable(y)))))
+    #     index_div_2 = prim.FloorDiv(index, 2)
+    #     index_div_2 = prim.Product((index_div_2, k-1))
+    #     modified_index = prim.Sum((modified_index, index_div_2))
+    #     if local_dimension() > 2:
+    #         modified_index = prim.Sum((modified_index, prim.Variable(z)))
 
-    modified_index = prim.Sum((prim.Variable(x), index))
-    if local_dimension()>1:
-        modified_index = prim.Sum((modified_index, prim.Product(((k+1), prim.Variable(y)))))
-    index_div_2 = prim.FloorDiv(index, 2)
-    index_div_2 = prim.Product((index_div_2, k-1))
-    modified_index = prim.Sum((modified_index, index_div_2))
+    modified_index = prim.Sum((inames, tensor_index_to_sequential_index(sublem_inames, get_option("number_of_blocks")+1)))
 
     return modified_index
 

python/dune/perftool/pdelab/argument.py

@@ -139,6 +139,9 @@ def pymbolic_coefficient(container, lfs, index):
     if isinstance(lfs, str):
         valuearg(lfs, dtype=NumpyType("str"))
 
+    if isinstance(index, tuple):
+        index = index[0]
+
     # If the LFS is not yet a pymbolic expression, make it one
     from pymbolic.primitives import Expression
     if not isinstance(lfs, Expression):

python/dune/perftool/pdelab/basis.py

@@ -14,6 +14,7 @@ from dune.perftool.pdelab.spaces import (lfs_child,
                                          name_lfs,
                                          name_lfs_bound,
                                          type_gfs,
+                                         lfs_inames
                                          )
 from dune.perftool.pdelab.geometry import (component_iname,
                                            world_dimension,
@@ -79,7 +80,6 @@ def declare_cache_temporary(element, restriction, which):
     return decl
 
 
-@backend(interface="evaluate_basis")
 @kernel_cached
 def evaluate_basis(leaf_element, name, restriction):
     lfs = name_leaf_lfs(leaf_element, restriction)
@@ -97,17 +97,17 @@ def evaluate_basis(leaf_element, name, restriction):
                 )
 
 
+@backend(interface="pymbolic_basis")
 def pymbolic_basis(leaf_element, restriction, number, context=''):
     assert leaf_element.num_sub_elements() == 0
     name = "phi_{}".format(FEM_name_mangling(leaf_element))
     name = restricted_name(name, restriction)
-    get_backend("evaluate_basis", selector=option_switch("blockstructured"))(leaf_element, name, restriction)
-    iname = get_backend("lfs_inames", selector=option_switch("blockstructured"))(leaf_element, restriction, number, context=context)[0]
+    evaluate_basis(leaf_element, name, restriction)
+    iname = lfs_inames(leaf_element, restriction, number, context=context)[0]
 
     return Subscript(Variable(name), (Variable(iname), ))
 
 
-@backend(interface="evaluate_grad")
 @kernel_cached
 def evaluate_reference_gradient(leaf_element, name, restriction):
     lfs = name_leaf_lfs(leaf_element, restriction)
@@ -125,12 +125,13 @@ def evaluate_reference_gradient(leaf_element, name, restriction):
                 )
 
 
+@backend(interface="pymbolic_reference_gradient")
 def pymbolic_reference_gradient(leaf_element, restriction, number, context=''):
     assert leaf_element.num_sub_elements() == 0
     name = "js_{}".format(FEM_name_mangling(leaf_element))
     name = restricted_name(name, restriction)
-    get_backend("evaluate_grad", selector=option_switch("blockstructured"))(leaf_element, name, restriction)
-    iname = get_backend("lfs_inames", selector=option_switch("blockstructured"))(leaf_element, restriction, number, context=context)[0]
+    evaluate_reference_gradient(leaf_element, name, restriction)
+    iname = lfs_inames(leaf_element, restriction, number, context=context)[0]
 
     return Subscript(Variable(name), (Variable(iname), 0))
 
@@ -199,7 +200,7 @@ def evaluate_coefficient_gradient(element, name, container, restriction, tree_pa
     # and proceed to call the necessary generator functions
     temporary_variable(name, shape=shape, shape_impl=shape_impl)
     lfs = name_lfs(element, restriction, tree_path)
-    basis = pymbolic_reference_gradient(leaf_element, restriction, 0, context='trialgrad')
+    basis = get_backend("pymbolic_reference_gradient", selector=option_switch("blockstructured"))(leaf_element, restriction, 0, context='trialgrad')
     index, _ = get_pymbolic_indices(basis)
 
     if isinstance(sub_element, (VectorElement, TensorElement)):

python/dune/perftool/pdelab/localoperator.py

@@ -243,11 +243,13 @@ def determine_accumulation_space(expr, number, measure, idims=None):
         lfs = lfs_child(lfs, idims, shape=subel.value_shape(), symmetry=subel.symmetry())
         subel = subel.sub_elements()[0]
 
-    iname = get_backend("lfs_inames", selector=option_switch("blockstructured"))(subel, ma.restriction, count=number)[0]
-    lfsi = Variable(iname)
+    iname = get_backend("lfs_inames", selector=option_switch("blockstructured"))(subel, ma.restriction, count=number)
     if get_option("blockstructured"):
         from dune.perftool.blockstructured.tools import micro_index_to_macro_index
-        lfsi = micro_index_to_macro_index(lfsi)
+        lfsi = micro_index_to_macro_index(iname)
+    else:
+        lfsi = Variable(iname[0])
+
 
     # If the LFS is not yet a pymbolic expression, make it one
     from pymbolic.primitives import Expression

python/dune/perftool/tools.py

@@ -23,8 +23,10 @@ def get_pymbolic_indices(expr):
 
     def ind(i):
         if isinstance(i, int):
-            return i
-        return get_pymbolic_basename(i)
+            return (i,)
+        if isinstance(i, prim.Sum) or isinstance(i, prim.Product):
+            return sum((ind(child) for child in i.children if isinstance(child, prim.Expression)), ())
+        return (get_pymbolic_basename(i),)
 
     if not isinstance(expr.index, tuple):
         return (ind(expr.index),)