From 7c0a1024136d55a84172dcc3ec9fb64c21fa6989 Mon Sep 17 00:00:00 2001 From: Connor Ward Date: Wed, 18 Dec 2024 09:35:06 +0000 Subject: [PATCH] Control logging of local kernels using a compiler parameter --- firedrake/interpolation.py | 8 +++----- firedrake/preconditioners/pmg.py | 2 +- firedrake/slate/slac/compiler.py | 2 +- firedrake/tsfc_interface.py | 6 +++--- tsfc/driver.py | 15 ++++++--------- tsfc/parameters.py | 3 +++ 6 files changed, 17 insertions(+), 19 deletions(-) diff --git a/firedrake/interpolation.py b/firedrake/interpolation.py index e8555ba7dd..0a598ba34b 100644 --- a/firedrake/interpolation.py +++ b/firedrake/interpolation.py @@ -1092,8 +1092,7 @@ def _interpolator(V, tensor, expr, subset, arguments, access, bcs=None): # interpolation) we have to pass the finat element we construct # here. Ideally we would only pass the UFL element through. kernel = compile_expression(cell_set.comm, expr, to_element, V.ufl_element(), - domain=source_mesh, parameters=parameters, - log=PETSc.Log.isActive()) + domain=source_mesh, parameters=parameters) ast = kernel.ast oriented = kernel.oriented needs_cell_sizes = kernel.needs_cell_sizes @@ -1221,10 +1220,9 @@ def _interpolator(V, tensor, expr, subset, arguments, access, bcs=None): f"firedrake-tsfc-expression-kernel-cache-uid{os.getuid()}") -def _compile_expression_key(comm, expr, to_element, ufl_element, domain, parameters, log): +def _compile_expression_key(comm, expr, to_element, ufl_element, domain, parameters): """Generate a cache key suitable for :func:`tsfc.compile_expression_dual_evaluation`.""" - key = hash_expr(expr), hash(ufl_element), utils.tuplify(parameters), log - return key + return (hash_expr(expr), hash(ufl_element), utils.tuplify(parameters)) @memory_and_disk_cache( diff --git a/firedrake/preconditioners/pmg.py b/firedrake/preconditioners/pmg.py index 251f585cbe..87524336ee 100644 --- a/firedrake/preconditioners/pmg.py +++ b/firedrake/preconditioners/pmg.py @@ -530,7 +530,7 @@ def coarsen_bc_value(self, bc, cV): 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(), log=PETSc.Log.isActive()) + kernel = compile_expression_dual_evaluation(expr, to_element, Vf.ufl_element()) coefficients = extract_numbered_coefficients(expr, kernel.coefficient_numbers) if kernel.needs_external_coords: coefficients = [Vf.mesh().coordinates] + coefficients diff --git a/firedrake/slate/slac/compiler.py b/firedrake/slate/slac/compiler.py index 7e1b14281c..d9e7bced02 100644 --- a/firedrake/slate/slac/compiler.py +++ b/firedrake/slate/slac/compiler.py @@ -237,4 +237,4 @@ def gem_to_loopy(gem_expr, var2terminal, scalar_type): # Part B: impero_c to loopy output_arg = OutputKernelArg(output_loopy_arg) - return generate_loopy(impero_c, args, scalar_type, "slate_loopy", [], log=PETSc.Log.isActive()), output_arg + return generate_loopy(impero_c, args, scalar_type, "slate_loopy", []), output_arg diff --git a/firedrake/tsfc_interface.py b/firedrake/tsfc_interface.py index a4a57ae0cb..ba10d79507 100644 --- a/firedrake/tsfc_interface.py +++ b/firedrake/tsfc_interface.py @@ -53,8 +53,8 @@ ) -def tsfc_compile_form_hashkey(form, prefix, parameters, interface, diagonal, log): - # Drop prefix as it's only used for naming and log +def tsfc_compile_form_hashkey(form, prefix, parameters, interface, diagonal): + # Drop prefix as it's only used for naming return default_parallel_hashkey(form.signature(), prefix, parameters, interface, diagonal) @@ -94,7 +94,7 @@ def __init__( """ tree = tsfc_compile_form(form, prefix=name, parameters=parameters, interface=interface, - diagonal=diagonal, log=PETSc.Log.isActive()) + diagonal=diagonal) kernels = [] for kernel in tree: # Individual kernels do not have to use all of the coefficients diff --git a/tsfc/driver.py b/tsfc/driver.py index 6e3c3baaf3..38a780a0ad 100644 --- a/tsfc/driver.py +++ b/tsfc/driver.py @@ -47,14 +47,13 @@ """ -def compile_form(form, prefix="form", parameters=None, interface=None, diagonal=False, log=False): +def compile_form(form, prefix="form", parameters=None, interface=None, diagonal=False): """Compiles a UFL form into a set of assembly kernels. :arg form: UFL form :arg prefix: kernel name will start with this string :arg parameters: parameters object :arg diagonal: Are we building a kernel for the diagonal of a rank-2 element tensor? - :arg log: bool if the Kernel should be profiled with Log events :returns: list of kernels """ cpu_time = time.time() @@ -71,7 +70,7 @@ def compile_form(form, prefix="form", parameters=None, interface=None, diagonal= kernels = [] for integral_data in fd.integral_data: start = time.time() - kernel = compile_integral(integral_data, fd, prefix, parameters, interface=interface, diagonal=diagonal, log=log) + kernel = compile_integral(integral_data, fd, prefix, parameters, interface=interface, diagonal=diagonal) if kernel is not None: kernels.append(kernel) logger.info(GREEN % "compile_integral finished in %g seconds.", time.time() - start) @@ -80,7 +79,7 @@ def compile_form(form, prefix="form", parameters=None, interface=None, diagonal= return kernels -def compile_integral(integral_data, form_data, prefix, parameters, interface, *, diagonal=False, log=False): +def compile_integral(integral_data, form_data, prefix, parameters, interface, *, diagonal=False): """Compiles a UFL integral into an assembly kernel. :arg integral_data: UFL integral data @@ -89,7 +88,6 @@ def compile_integral(integral_data, form_data, prefix, parameters, interface, *, :arg parameters: parameters object :arg interface: backend module for the kernel interface :arg diagonal: Are we building a kernel for the diagonal of a rank-2 element tensor? - :arg log: bool if the Kernel should be profiled with Log events :returns: a kernel constructed by the kernel interface """ parameters = preprocess_parameters(parameters) @@ -137,7 +135,7 @@ def compile_integral(integral_data, form_data, prefix, parameters, interface, *, integrand_exprs = builder.compile_integrand(integrand, params, ctx) integral_exprs = builder.construct_integrals(integrand_exprs, params) builder.stash_integrals(integral_exprs, params, ctx) - return builder.construct_kernel(kernel_name, ctx, log) + return builder.construct_kernel(kernel_name, ctx, parameters["add_petsc_events"]) def preprocess_parameters(parameters): @@ -157,7 +155,7 @@ def preprocess_parameters(parameters): def compile_expression_dual_evaluation(expression, to_element, ufl_element, *, domain=None, interface=None, - parameters=None, log=False): + parameters=None): """Compile a UFL expression to be evaluated against a compile-time known reference element's dual basis. Useful for interpolating UFL expressions into e.g. N1curl spaces. @@ -168,7 +166,6 @@ def compile_expression_dual_evaluation(expression, to_element, ufl_element, *, :arg domain: optional UFL domain the expression is defined on (required when expression contains no domain). :arg interface: backend module for the kernel interface :arg parameters: parameters object - :arg log: bool if the Kernel should be profiled with Log events :returns: Loopy-based ExpressionKernel object. """ if parameters is None: @@ -267,7 +264,7 @@ def compile_expression_dual_evaluation(expression, to_element, ufl_element, *, builder.register_requirements([evaluation]) builder.set_output(return_var) # Build kernel tuple - return builder.construct_kernel(impero_c, index_names, needs_external_coords, log=log) + return builder.construct_kernel(impero_c, index_names, needs_external_coords, parameters["add_petsc_events"]) class DualEvaluationCallable(object): diff --git a/tsfc/parameters.py b/tsfc/parameters.py index 1277713ad5..af44ce0cd4 100644 --- a/tsfc/parameters.py +++ b/tsfc/parameters.py @@ -20,6 +20,9 @@ # So that tests pass (needs to match scalar_type) "scalar_type_c": "double", + + # Whether to wrap the generated kernels in a PETSc event + "add_petsc_events": False, }