Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Control logging of local kernels using a compiler parameter #3937

Merged
merged 1 commit into from
Dec 18, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 3 additions & 5 deletions firedrake/interpolation.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand Down
2 changes: 1 addition & 1 deletion firedrake/preconditioners/pmg.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion firedrake/slate/slac/compiler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
6 changes: 3 additions & 3 deletions firedrake/tsfc_interface.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)


Expand Down Expand Up @@ -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
Expand Down
15 changes: 6 additions & 9 deletions tsfc/driver.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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)
Expand All @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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):
Expand All @@ -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.
Expand All @@ -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:
Expand Down Expand Up @@ -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):
Expand Down
3 changes: 3 additions & 0 deletions tsfc/parameters.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}


Expand Down
Loading