From fddc94d3ab1382883d6fd4846bee518401cc9020 Mon Sep 17 00:00:00 2001
From: Pablo Brubeck <brubeck@protonmail.com>
Date: Fri, 13 Dec 2024 15:58:28 +0000
Subject: [PATCH] TSFC: Cache basis evaluation (#3921)

* TSFC: cache basis evaluation

* TSFC: cache QuadratureRule
---
 tsfc/fem.py                     | 24 +++++++++++++++++-------
 tsfc/kernel_interface/common.py | 14 ++++++--------
 2 files changed, 23 insertions(+), 15 deletions(-)

diff --git a/tsfc/fem.py b/tsfc/fem.py
index abc8bc7cb8..99251ed0c6 100644
--- a/tsfc/fem.py
+++ b/tsfc/fem.py
@@ -3,7 +3,7 @@
 
 import collections
 import itertools
-from functools import singledispatch
+from functools import cached_property, singledispatch
 
 import gem
 import numpy
@@ -18,7 +18,6 @@
 from gem.node import traversal
 from gem.optimise import constant_fold_zero, ffc_rounding
 from gem.unconcatenate import unconcatenate
-from gem.utils import cached_property
 from ufl.classes import (Argument, CellCoordinate, CellEdgeVectors,
                          CellFacetJacobian, CellOrientation, CellOrigin,
                          CellVertices, CellVolume, Coefficient, FacetArea,
@@ -42,6 +41,8 @@
                             TSFCConstantMixin, entity_avg, one_times,
                             preprocess_expression, simplify_abs)
 
+from pyop2.caching import serial_cache
+
 
 class ContextBase(ProxyKernelInterface):
     """Common UFL -> GEM translation context."""
@@ -262,6 +263,18 @@ def needs_coordinate_mapping(element):
         return isinstance(create_element(element), NeedsCoordinateMappingElement)
 
 
+@serial_cache(hashkey=lambda *args: args)
+def get_quadrature_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 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."""
 
@@ -272,13 +285,9 @@ 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)
+        return get_quadrature_rule(self.fiat_cell, self.integration_dim, self.quadrature_degree, "default")
 
     @cached_property
     def point_set(self):
@@ -296,6 +305,7 @@ def point_expr(self):
     def weight_expr(self):
         return self.quadrature_rule.weight_expression
 
+    @serial_cache(hashkey=make_basis_evaluation_key)
     def basis_evaluation(self, finat_element, mt, entity_id):
         return finat_element.basis_evaluation(mt.local_derivatives,
                                               self.point_set,
diff --git a/tsfc/kernel_interface/common.py b/tsfc/kernel_interface/common.py
index df7e879f09..f757af951f 100644
--- a/tsfc/kernel_interface/common.py
+++ b/tsfc/kernel_interface/common.py
@@ -1,7 +1,7 @@
 import collections
 import operator
 import string
-from functools import reduce
+from functools import cached_property, reduce
 from itertools import chain, product
 
 import gem
@@ -9,11 +9,10 @@
 import numpy
 from FIAT.reference_element import TensorProductCell
 from finat.cell_tools import max_complex
-from finat.quadrature import AbstractQuadratureRule, make_quadrature
+from finat.quadrature import AbstractQuadratureRule
 from gem.node import traversal
 from gem.optimise import constant_fold_zero
 from gem.optimise import remove_componenttensors as prune
-from gem.utils import cached_property
 from numpy import asarray
 from tsfc import fem, ufl_utils
 from tsfc.finatinterface import as_fiat_cell, create_element
@@ -137,7 +136,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()
@@ -320,8 +319,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 = fem.get_quadrature_rule(fiat_cell, integration_dim, quadrature_degree, scheme)
         params["quadrature_rule"] = quad_rule
 
     if not isinstance(quad_rule, AbstractQuadratureRule):
@@ -330,8 +328,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