Do not introduce too much shape for reference gradient

python/dune/perftool/sumfact/__init__.py

@@ -29,7 +29,9 @@ class SumFactInterface(PDELabInterface):
         return pymbolic_basis(element, restriction, number)
 
     def pymbolic_reference_gradient(self, element, restriction, number):
-        return pymbolic_reference_gradient(element, restriction, number)
+        ret, indices = pymbolic_reference_gradient(element, restriction, number, self.visitor.indices)
+        self.visitor.indices = indices
+        return ret
 
     def pymbolic_trialfunction_gradient(self, element, restriction, component):
         ret, indices = pymbolic_coefficient_gradient(element, restriction, component, name_coefficientcontainer, self.visitor.indices)

python/dune/perftool/sumfact/basis.py

@@ -271,9 +271,9 @@ def pymbolic_basis(element, restriction, number):
 
 @backend(interface="evaluate_grad")
 @kernel_cached
-def evaluate_reference_gradient(element, name, restriction):
+def evaluate_reference_gradient(element, name, restriction, index):
     dim = world_dimension()
-    temporary_variable(name, shape=(dim,))
+    temporary_variable(name, shape=())
     quad_inames = quadrature_inames()
     inames = lfs_inames(element, restriction)
     facedir = get_facedir(restriction)
@@ -287,35 +287,34 @@ def evaluate_reference_gradient(element, name, restriction):
         quadinamemapping[i] = quad_inames[d]
         i = i + 1
 
-    for d in range(dim):
-        prod = []
-        for i in range(dim):
-            if i != facedir:
-                prod.append(BasisTabulationMatrix(derivative=d == i).pymbolic((prim.Variable(quadinamemapping[i]), prim.Variable(inames[i]))))
-        if facedir is not None:
-            # TODO: Does not work for systems!
-            degree = polynomial_degree()
-
-            facemod = get_facemod(restriction)
-            prod.append(prim.Call(PolynomialLookup(name_polynomials(degree), facedir == d),
-                                  (prim.Variable(inames[facedir]), facemod)),)
+    prod = []
+    for i in range(dim):
+        if i != facedir:
+            prod.append(BasisTabulationMatrix(derivative=index == i).pymbolic((prim.Variable(quadinamemapping[i]), prim.Variable(inames[i]))))
+    if facedir is not None:
+        # TODO: Does not work for systems!
+        degree = polynomial_degree()
 
-        assignee = prim.Subscript(prim.Variable(name), (d,))
-        expression = prim.Product(tuple(prod))
+        facemod = get_facemod(restriction)
+        prod.append(prim.Call(PolynomialLookup(name_polynomials(degree), facedir == d),
+                              (prim.Variable(inames[facedir]), facemod)),)
 
-        instruction(assignee=assignee,
-                    expression=expression,
-                    forced_iname_deps=frozenset(quad_inames + inames),
-                    forced_iname_deps_is_final=True,
-                    )
+    instruction(assignee=prim.Variable(name),
+                expression=prim.Product(tuple(prod)),
+                forced_iname_deps=frozenset(quad_inames + inames),
+                forced_iname_deps_is_final=True,
+                )
 
 
-def pymbolic_reference_gradient(element, restriction, number):
+def pymbolic_reference_gradient(element, restriction, number, indices):
     assert number == 1
+    assert len(indices) == 1
+    index, = indices
 
     # TODO: Change name?
     name = "js_{}".format(FEM_name_mangling(element))
     name = restricted_name(name, restriction)
-    evaluate_reference_gradient(element, name, restriction)
+    name = "{}_{}".format(name, index)
+    evaluate_reference_gradient(element, name, restriction, index)
 
-    return prim.Variable(name)
+    return prim.Variable(name), None

python/dune/perftool/ufl/visitor.py

@@ -69,11 +69,6 @@ class UFL2LoopyVisitor(ModifiedTerminalTracker):
         element = select_subelement(o.ufl_element(), self.component)
         leaf_element = element
 
-        indices = self.indices
-        if indices is None:
-            indices = ()
-        self.indices = None
-
         # Now treat the case of this being a vector finite element
         if element.num_sub_elements() > 0:
             # I cannot handle general mixed elements here...
@@ -81,7 +76,7 @@ class UFL2LoopyVisitor(ModifiedTerminalTracker):
 
             # If this is a vector element, we need add an additional accumulation loop iname
             shape = len(element.value_shape())
-            indices = indices[shape:]
+            self.indices = self.indices[shape:]
             for i in range(len(element.value_shape())):
                 if self.interface.dimension_iname(context='arg', count=i) not in self.inames:
                     self.inames = self.inames + (self.interface.dimension_iname(context='arg', count=i),)
@@ -99,9 +94,9 @@ class UFL2LoopyVisitor(ModifiedTerminalTracker):
                 self.inames = self.inames + (iname,)
 
         if self.reference_grad:
-            return maybe_wrap_subscript(self.interface.pymbolic_reference_gradient(leaf_element, restriction, o.number()), indices)
+            return self.interface.pymbolic_reference_gradient(leaf_element, restriction, o.number())
         else:
-            return maybe_wrap_subscript(self.interface.pymbolic_basis(leaf_element, restriction, o.number()), indices)
+            return self.interface.pymbolic_basis(leaf_element, restriction, o.number())
 
     def coefficient(self, o):
         # Do something different for trial function and coefficients from jacobian apply