diff --git a/tsfc/fem.py b/tsfc/fem.py index 1d5b1ac392..00cf3435ec 100644 --- a/tsfc/fem.py +++ b/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 (finat_element, mt.local_derivatives, ctx.point_set, ctx.integration_dim, entity_id, coordinate_element, mt.restriction) + + 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, diff --git a/tsfc/kernel_interface/common.py b/tsfc/kernel_interface/common.py index 5cb2096961..5081ad0c96 100644 --- a/tsfc/kernel_interface/common.py +++ b/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