TSFC: cache QuadratureRule

tsfc/fem.py

@@ -14,7 +14,6 @@
 from finat.physically_mapped import (NeedsCoordinateMappingElement,
                                      PhysicalGeometry)
 from finat.point_set import PointSet, PointSingleton
-from finat.quadrature import make_quadrature
 from gem.node import traversal
 from gem.optimise import constant_fold_zero, ffc_rounding
 from gem.unconcatenate import unconcatenate
@@ -263,6 +262,12 @@ def needs_coordinate_mapping(element):
         return isinstance(create_element(element), NeedsCoordinateMappingElement)
 
 
+def make_basis_evaluation_key(ctx, finat_element, mt, entity_id):
+    domain = extract_unique_domain(mt.terminal)
+    coordinate_element = domain.ufl_coordinate_element()
+    return (ctx.point_set, finat_element, mt.local_derivatives, coordinate_element, mt.restriction, ctx.integration_dim, entity_id)
+
+
 class PointSetContext(ContextBase):
     """Context for compile-time known evaluation points."""
 
@@ -273,14 +278,6 @@ class PointSetContext(ContextBase):
         'weight_expr',
     )
 
-    @cached_property
-    def integration_cell(self):
-        return self.fiat_cell.construct_subelement(self.integration_dim)
-
-    @cached_property
-    def quadrature_rule(self):
-        return make_quadrature(self.integration_cell, self.quadrature_degree)
-
     @cached_property
     def point_set(self):
         return self.quadrature_rule.point_set
@@ -297,11 +294,6 @@ def point_expr(self):
     def weight_expr(self):
         return self.quadrature_rule.weight_expression
 
-    def make_basis_evaluation_key(self, finat_element, mt, entity_id):
-        domain = extract_unique_domain(mt.terminal)
-        coordinate_element = domain.ufl_coordinate_element()
-        return (self, finat_element, mt.local_derivatives, coordinate_element, mt.restriction, entity_id)
-
     @serial_cache(hashkey=make_basis_evaluation_key)
     def basis_evaluation(self, finat_element, mt, entity_id):
         return finat_element.basis_evaluation(mt.local_derivatives,
@@ -319,6 +311,7 @@ class GemPointContext(ContextBase):
         'weight_expr',
     )
 
+    @serial_cache(hashkey=make_basis_evaluation_key)
     def basis_evaluation(self, finat_element, mt, entity_id):
         return finat_element.point_evaluation(mt.local_derivatives,
                                               self.point_expr,

tsfc/kernel_interface/common.py

@@ -20,6 +20,8 @@
 from tsfc.logging import logger
 from ufl.utils.sequences import max_degree
 
+from pyop2.caching import serial_cache
+
 
 class KernelBuilderBase(KernelInterface):
     """Helper class for building local assembly kernels."""
@@ -136,7 +138,7 @@ def compile_integrand(self, integrand, params, ctx):
             integrand = ufl_utils.split_coefficients(integrand, self.coefficient_split)
         # Compile: ufl -> gem
         info = self.integral_data_info
-        functions = list(info.arguments) + [self.coordinate(info.domain)] + list(info.coefficients)
+        functions = [*info.arguments, self.coordinate(info.domain), *info.coefficients]
         set_quad_rule(params, info.domain.ufl_cell(), info.integral_type, functions)
         quad_rule = params["quadrature_rule"]
         config = self.fem_config()
@@ -294,6 +296,12 @@ def create_context(self):
                 'mode_irs': collections.OrderedDict()}
 
 
+@serial_cache(hashkey=lambda *args: args)
+def cache_quad_rule(fiat_cell, integration_dim, quadrature_degree, scheme):
+    integration_cell = fiat_cell.construct_subcomplex(integration_dim)
+    return make_quadrature(integration_cell, quadrature_degree, scheme=scheme)
+
+
 def set_quad_rule(params, cell, integral_type, functions):
     # Check if the integral has a quad degree attached, otherwise use
     # the estimated polynomial degree attached by compute_form_data
@@ -319,8 +327,7 @@ def set_quad_rule(params, cell, integral_type, functions):
         fiat_cell = max_complex(fiat_cells)
 
         integration_dim, _ = lower_integral_type(fiat_cell, integral_type)
-        integration_cell = fiat_cell.construct_subcomplex(integration_dim)
-        quad_rule = make_quadrature(integration_cell, quadrature_degree, scheme=scheme)
+        quad_rule = cache_quad_rule(fiat_cell, integration_dim, quadrature_degree, scheme)
         params["quadrature_rule"] = quad_rule
 
     if not isinstance(quad_rule, AbstractQuadratureRule):
@@ -329,8 +336,8 @@ def set_quad_rule(params, cell, integral_type, functions):
 
 
 def get_index_ordering(quadrature_indices, return_variables):
-    split_argument_indices = tuple(chain(*[var.index_ordering()
-                                           for var in return_variables]))
+    split_argument_indices = tuple(chain(*(var.index_ordering()
+                                           for var in return_variables)))
     return tuple(quadrature_indices) + split_argument_indices