From 2fd0fb2ac8c999ff89e9a3874857d072f9e8f12f Mon Sep 17 00:00:00 2001 From: Connor Ward Date: Wed, 16 Feb 2022 11:37:23 +0000 Subject: [PATCH 1/3] Replace ExpressionKernel.coefficients with coefficient numbers This should enable caching of ExpressionKernels since they are no longer 'heavy'. --- firedrake/interpolation.py | 9 +++++---- firedrake/preconditioners/pmg.py | 6 +++--- firedrake/tsfc_interface.py | 18 ++++++++++++++++++ 3 files changed, 26 insertions(+), 7 deletions(-) diff --git a/firedrake/interpolation.py b/firedrake/interpolation.py index 486a03872d..b23f255bd1 100644 --- a/firedrake/interpolation.py +++ b/firedrake/interpolation.py @@ -3,7 +3,7 @@ import FIAT import ufl -from ufl.algorithms import extract_arguments +from ufl.algorithms import extract_arguments, extract_coefficients from pyop2 import op2 @@ -14,7 +14,7 @@ import finat import firedrake -from firedrake import utils +from firedrake import tsfc_interface, utils from firedrake.adjoint import annotate_interpolate from firedrake.petsc import PETSc @@ -286,7 +286,7 @@ def _interpolator(V, tensor, expr, subset, arguments, access): ast = kernel.ast oriented = kernel.oriented needs_cell_sizes = kernel.needs_cell_sizes - coefficients = kernel.coefficients + coefficient_numbers = kernel.coefficient_numbers first_coeff_fake_coords = kernel.first_coefficient_fake_coords name = kernel.name kernel = op2.Kernel(ast, name, requires_zeroed_output_arguments=True, @@ -297,9 +297,10 @@ def _interpolator(V, tensor, expr, subset, arguments, access): cell_set = subset parloop_args = [kernel, cell_set] + coefficients = tsfc_interface.extract_numbered_coefficients(expr, coefficient_numbers) if first_coeff_fake_coords: # Replace with real source mesh coordinates - coefficients[0] = source_mesh.coordinates + coefficients = [source_mesh.coordinates] + coefficients if target_mesh is not source_mesh: # NOTE: TSFC will sometimes drop run-time arguments in generated diff --git a/firedrake/preconditioners/pmg.py b/firedrake/preconditioners/pmg.py index a814f4cf0d..4fc3d4a2df 100644 --- a/firedrake/preconditioners/pmg.py +++ b/firedrake/preconditioners/pmg.py @@ -6,6 +6,7 @@ add_hook, get_parent, push_parent, pop_parent, get_function_space, set_function_space) from firedrake.solving_utils import _SNESContext +from firedrake.tsfc_interface import extract_numbered_coefficients from firedrake.utils import ScalarType_c, IntType_c from firedrake.petsc import PETSc import firedrake @@ -500,10 +501,9 @@ def prolongation_transfer_kernel_action(Vf, expr): from tsfc.finatinterface import create_element to_element = create_element(Vf.ufl_element()) kernel = compile_expression_dual_evaluation(expr, to_element, Vf.ufl_element()) - coefficients = kernel.coefficients + coefficients = extract_numbered_coefficients(expr, kernel.coefficient_numbers) if kernel.first_coefficient_fake_coords: - target_mesh = Vf.ufl_domain() - coefficients[0] = target_mesh.coordinates + coefficients = [Vf.ufl_domain().coordinates] + coefficients return op2.Kernel(kernel.ast, kernel.name, requires_zeroed_output_arguments=True, diff --git a/firedrake/tsfc_interface.py b/firedrake/tsfc_interface.py index 49d34dd020..930305c5ac 100644 --- a/firedrake/tsfc_interface.py +++ b/firedrake/tsfc_interface.py @@ -304,3 +304,21 @@ def as_pyop2_local_kernel(ast, name, nargs, access=op2.INC, **kwargs): accesses = tuple([access] + [op2.READ]*(nargs-1)) return op2.Kernel(ast, name, accesses=accesses, requires_zeroed_output_arguments=True, **kwargs) + + +def extract_numbered_coefficients(expr, numbers): + """Return expression coefficients specified by a numbering. + + :arg expr: A UFL expression. + :arg numbers: Iterable of indices used for selecting the correct coefficients + from ``expr``. + :returns: A list of UFL coefficients. + """ + orig_coefficients = ufl.algorithms.extract_coefficients(expr) + coefficients = [] + for coeff in (orig_coefficients[i] for i in numbers): + if type(coeff.ufl_element()) == ufl.MixedElement: + coefficients.extend(coeff.split()) + else: + coefficients.append(coeff) + return coefficients From c2eee14b45e1f5eb2a9739fa88cb62059c70bd32 Mon Sep 17 00:00:00 2001 From: Connor Ward Date: Fri, 1 Apr 2022 10:53:30 +0100 Subject: [PATCH 2/3] Rename first_coefficient_fake_coords to needs_external_coords --- firedrake/interpolation.py | 5 ++--- firedrake/preconditioners/pmg.py | 2 +- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/firedrake/interpolation.py b/firedrake/interpolation.py index b23f255bd1..306e9d30cc 100644 --- a/firedrake/interpolation.py +++ b/firedrake/interpolation.py @@ -287,7 +287,7 @@ def _interpolator(V, tensor, expr, subset, arguments, access): oriented = kernel.oriented needs_cell_sizes = kernel.needs_cell_sizes coefficient_numbers = kernel.coefficient_numbers - first_coeff_fake_coords = kernel.first_coefficient_fake_coords + needs_external_coords = kernel.needs_external_coords name = kernel.name kernel = op2.Kernel(ast, name, requires_zeroed_output_arguments=True, flop_count=kernel.flop_count) @@ -298,8 +298,7 @@ def _interpolator(V, tensor, expr, subset, arguments, access): parloop_args = [kernel, cell_set] coefficients = tsfc_interface.extract_numbered_coefficients(expr, coefficient_numbers) - if first_coeff_fake_coords: - # Replace with real source mesh coordinates + if needs_external_coords: coefficients = [source_mesh.coordinates] + coefficients if target_mesh is not source_mesh: diff --git a/firedrake/preconditioners/pmg.py b/firedrake/preconditioners/pmg.py index 4fc3d4a2df..0ad931b930 100644 --- a/firedrake/preconditioners/pmg.py +++ b/firedrake/preconditioners/pmg.py @@ -502,7 +502,7 @@ def prolongation_transfer_kernel_action(Vf, expr): to_element = create_element(Vf.ufl_element()) kernel = compile_expression_dual_evaluation(expr, to_element, Vf.ufl_element()) coefficients = extract_numbered_coefficients(expr, kernel.coefficient_numbers) - if kernel.first_coefficient_fake_coords: + if kernel.needs_external_coords: coefficients = [Vf.ufl_domain().coordinates] + coefficients return op2.Kernel(kernel.ast, kernel.name, From 3d356da7bfcfa6aae144fe6200bb1535497981e8 Mon Sep 17 00:00:00 2001 From: Connor Ward Date: Wed, 16 Feb 2022 11:40:08 +0000 Subject: [PATCH 3/3] Add disk caching for interpolation kernels --- firedrake/interpolation.py | 51 ++++++++++++++++++++++++++++++++------ 1 file changed, 43 insertions(+), 8 deletions(-) diff --git a/firedrake/interpolation.py b/firedrake/interpolation.py index 306e9d30cc..3c6b706fd2 100644 --- a/firedrake/interpolation.py +++ b/firedrake/interpolation.py @@ -1,11 +1,15 @@ import numpy from functools import partial, singledispatch +import os +import tempfile import FIAT import ufl from ufl.algorithms import extract_arguments, extract_coefficients +from ufl.algorithms.signature import compute_expression_signature from pyop2 import op2 +from pyop2.caching import disk_cached from tsfc.finatinterface import create_element, as_fiat_cell from tsfc import compile_expression_dual_evaluation @@ -270,6 +274,11 @@ def _interpolator(V, tensor, expr, subset, arguments, access): rt_var_name = 'rt_X' to_element = rebuild(to_element, expr, rt_var_name) + cell_set = target_mesh.cell_set + if subset is not None: + assert subset.superset == cell_set + cell_set = subset + parameters = {} parameters['scalar_type'] = utils.ScalarType @@ -279,10 +288,8 @@ def _interpolator(V, tensor, expr, subset, arguments, access): # FIXME: for the runtime unknown point set (for cross-mesh # interpolation) we have to pass the finat element we construct # here. Ideally we would only pass the UFL element through. - kernel = compile_expression_dual_evaluation(expr, to_element, - V.ufl_element(), - domain=source_mesh, - parameters=parameters) + kernel = compile_expression(cell_set.comm, expr, to_element, V.ufl_element(), + domain=source_mesh, parameters=parameters) ast = kernel.ast oriented = kernel.oriented needs_cell_sizes = kernel.needs_cell_sizes @@ -291,10 +298,6 @@ def _interpolator(V, tensor, expr, subset, arguments, access): name = kernel.name kernel = op2.Kernel(ast, name, requires_zeroed_output_arguments=True, flop_count=kernel.flop_count) - cell_set = target_mesh.cell_set - if subset is not None: - assert subset.superset == cell_set - cell_set = subset parloop_args = [kernel, cell_set] coefficients = tsfc_interface.extract_numbered_coefficients(expr, coefficient_numbers) @@ -381,6 +384,27 @@ def _interpolator(V, tensor, expr, subset, arguments, access): return copyin + (parloop_compute_callable, ) + copyout +try: + _expr_cachedir = os.environ["FIREDRAKE_TSFC_KERNEL_CACHE_DIR"] +except KeyError: + _expr_cachedir = os.path.join(tempfile.gettempdir(), + f"firedrake-tsfc-expression-kernel-cache-uid{os.getuid()}") + + +def _compile_expression_key(comm, expr, to_element, ufl_element, domain, parameters): + """Generate a cache key suitable for :func:`tsfc.compile_expression_dual_evaluation`.""" + # Since the caching is collective, this function must return a 2-tuple of + # the form (comm, key) where comm is the communicator the cache is collective over. + # FIXME FInAT elements are not safely hashable so we ignore them here + key = _hash_expr(expr), hash(ufl_element), utils.tuplify(parameters) + return comm, key + + +@disk_cached({}, _expr_cachedir, key=_compile_expression_key, collective=True) +def compile_expression(comm, *args, **kwargs): + return compile_expression_dual_evaluation(*args, **kwargs) + + @singledispatch def rebuild(element, expr, rt_var_name): raise NotImplementedError(f"Cross mesh interpolation not implemented for a {element} element.") @@ -490,3 +514,14 @@ def __init__(self, glob): self.dat = glob self.cell_node_map = lambda *arguments: None self.ufl_domain = lambda: None + + +def _hash_expr(expr): + """Return a numbering-invariant hash of a UFL expression. + + :arg expr: A UFL expression. + :returns: A numbering-invariant hash for the expression. + """ + domain_numbering = {d: i for i, d in enumerate(ufl.domain.extract_domains(expr))} + coefficient_numbering = {c: i for i, c in enumerate(extract_coefficients(expr))} + return compute_expression_signature(expr, {**domain_numbering, **coefficient_numbering})