From 3905e1c83ea730222100d654e0296d2390e052f4 Mon Sep 17 00:00:00 2001 From: Matthew Walker <247964090+mattatmet@users.noreply.github.com> Date: Tue, 7 Apr 2026 11:55:52 +0100 Subject: [PATCH 1/4] Adding profiling function to psyclone_tools --- .../build/psyclone/psyclone_tools.py | 37 ++++++++++++++++++- 1 file changed, 36 insertions(+), 1 deletion(-) diff --git a/infrastructure/build/psyclone/psyclone_tools.py b/infrastructure/build/psyclone/psyclone_tools.py index 14060c3e2..47a74c587 100644 --- a/infrastructure/build/psyclone/psyclone_tools.py +++ b/infrastructure/build/psyclone/psyclone_tools.py @@ -13,13 +13,14 @@ from psyclone.domain.lfric import LFRicConstants from psyclone.psyGen import InvokeSchedule -from psyclone.psyir.nodes import Loop, Routine, Directive +from psyclone.psyir.nodes import Loop, Routine, Directive, Container from psyclone.transformations import ( Dynamo0p3ColourTrans, Dynamo0p3OMPLoopTrans, Dynamo0p3RedundantComputationTrans, OMPParallelTrans, ) +from psyclone.psyir.transformations import ProfileTrans # List of allowed 'setval_*' built-ins for redundant computation transformation SETVAL_BUILTINS = ["setval_c"] @@ -75,6 +76,7 @@ def colour_loops(psyir, enable_tiling=False): It creates the instance of `Dynamo0p3ColourTrans` only once. :param psyir: the PSyIR of the PSy-layer. + :param enable_tiling: a bool to enable tiling. Default False. :type psyir: :py:class:`psyclone.psyir.nodes.FileContainer` """ @@ -94,6 +96,39 @@ def colour_loops(psyir, enable_tiling=False): ): ctrans.apply(child, options={"tiling": enable_tiling}) +# ----------------------------------------------------------------------------- +def profile_loops(psyir,colours_only=True): + """ + Applies timing calipers to kernels during the psyclone build. The default + is to only profile coloured loops but colours_only can be set to False to + profile every instance of a coded kernel. + + :param psyir: the PSyIR of the PSy-layer. + :param colours_only: profile only the coloured kernels. Default True. + :type psyir: :py:class:`psyclone.psyir.nodes.FileContainer` + + """ + profile_trans = ProfileTrans() + leave_loops = ["cells_in_colour", + "tiles_in_colour", + "cells_in_tile"] + + # Loop over all the InvokeSchedule in the PSyIR object + for subroutine in psyir.walk(InvokeSchedule): + # Add timing calipers to coloured loops. This should be done + # before the application of the openmp transformation. + count = 0 + for loop in subroutine.loops(): + if not loop.coded_kernels(): + continue + # Insert profiler calls before loop over colours + if (not colours_only and not loop.loop_type in leave_loops) or loop.loop_type == "colours": + k_name = loop.ancestor(InvokeSchedule).coded_kernels()[count].name + invoke_name = loop.ancestor(InvokeSchedule).invoke.name + file_name = loop.ancestor(Container).name + options = {"region_name": (file_name,invoke_name + ":" + k_name + "_k" + str(count))} + profile_trans.apply(loop,options=options) + count += 1 # ----------------------------------------------------------------------------- def openmp_parallelise_loops(psyir): From da0d83adc32d72ab36032ca27fde6da383168b53 Mon Sep 17 00:00:00 2001 From: Matthew Walker <247964090+mattatmet@users.noreply.github.com> Date: Tue, 21 Apr 2026 10:19:50 +0100 Subject: [PATCH 2/4] Adding my name to contributors list --- CONTRIBUTORS.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index 92ef52a5d..528bcd52e 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -25,5 +25,5 @@ | Pierre-siddall | Pierre Siddall | Met Office | 2026-01-29 | | mo-lucy-gordon | Lucy Gordon | Met Office | 2026-03-18 | | shreybh1 | Shrey Bhardwaj | Met Office | 2026-03-26 | - +| mattatmet | Matthew Walker | Met Office | 2026-04-21 | From 078febd9dcf06e0e2f0db61d77660c19b12352cc Mon Sep 17 00:00:00 2001 From: Matthew Walker <247964090+mattatmet@users.noreply.github.com> Date: Thu, 23 Apr 2026 14:58:52 +0100 Subject: [PATCH 3/4] Addition of type hints and order checks --- .../build/psyclone/psyclone_tools.py | 32 +++++++++++++++---- 1 file changed, 26 insertions(+), 6 deletions(-) diff --git a/infrastructure/build/psyclone/psyclone_tools.py b/infrastructure/build/psyclone/psyclone_tools.py index 47a74c587..7c671cbb5 100644 --- a/infrastructure/build/psyclone/psyclone_tools.py +++ b/infrastructure/build/psyclone/psyclone_tools.py @@ -13,12 +13,23 @@ from psyclone.domain.lfric import LFRicConstants from psyclone.psyGen import InvokeSchedule -from psyclone.psyir.nodes import Loop, Routine, Directive, Container +from psyclone.psyir.nodes import ( + Loop, + Routine, + Directive, + Container, + OMPParallelDirective, + OMPParallelDoDirective, + OMPDoDirective, + FileContainer, + ProfileNode +) from psyclone.transformations import ( Dynamo0p3ColourTrans, Dynamo0p3OMPLoopTrans, Dynamo0p3RedundantComputationTrans, OMPParallelTrans, + TransformationError ) from psyclone.psyir.transformations import ProfileTrans @@ -27,7 +38,7 @@ # ----------------------------------------------------------------------------- -def redundant_computation_setval(psyir): +def redundant_computation_setval(psyir: FileContainer): """ Applies the redundant computation transformation to loops over DoFs for the initialision built-ins, 'setval_*'. @@ -69,7 +80,7 @@ def redundant_computation_setval(psyir): # ----------------------------------------------------------------------------- -def colour_loops(psyir, enable_tiling=False): +def colour_loops(psyir: FileContainer, enable_tiling=False): """ Applies the colouring transformation to all applicable loops and optionally enables tiling. @@ -88,6 +99,10 @@ def colour_loops(psyir, enable_tiling=False): # Colour loops over cells unless they are on discontinuous # spaces or over DoFs for child in subroutine.children: + # Check if the profiling calipers have been added before the colouring. + if isinstance(child, ProfileNode): + raise TransformationError( + "Must apply colour_loops BEFORE profile_loops function in optimisation script.") if ( isinstance(child, Loop) and child.iteration_space.endswith("cell_column") @@ -97,7 +112,7 @@ def colour_loops(psyir, enable_tiling=False): ctrans.apply(child, options={"tiling": enable_tiling}) # ----------------------------------------------------------------------------- -def profile_loops(psyir,colours_only=True): +def profile_loops(psyir: FileContainer,colours_only=True): """ Applies timing calipers to kernels during the psyclone build. The default is to only profile coloured loops but colours_only can be set to False to @@ -123,6 +138,11 @@ def profile_loops(psyir,colours_only=True): continue # Insert profiler calls before loop over colours if (not colours_only and not loop.loop_type in leave_loops) or loop.loop_type == "colours": + # First check that the transformation is not being made inside an OMP region. + if loop.ancestor(OMPParallelDirective) or loop.ancestor(OMPParallelDoDirective) \ + or loop.ancestor(OMPDoDirective): + raise TransformationError( + "Must apply profile_loops BEFORE openmp_parellelise_loops function in optimisation script.") k_name = loop.ancestor(InvokeSchedule).coded_kernels()[count].name invoke_name = loop.ancestor(InvokeSchedule).invoke.name file_name = loop.ancestor(Container).name @@ -131,7 +151,7 @@ def profile_loops(psyir,colours_only=True): count += 1 # ----------------------------------------------------------------------------- -def openmp_parallelise_loops(psyir): +def openmp_parallelise_loops(psyir: FileContainer): """ Applies OpenMP Loop transformation to each applicable loop. @@ -155,7 +175,7 @@ def openmp_parallelise_loops(psyir): # ----------------------------------------------------------------------------- -def view_transformed_schedule(psyir): +def view_transformed_schedule(psyir: FileContainer): """ Provides view of transformed Invoke schedule in the PSy-layer. From 32cf0ad249647359d1d0da28c1f266c4706188f0 Mon Sep 17 00:00:00 2001 From: Matthew Walker <247964090+mattatmet@users.noreply.github.com> Date: Mon, 27 Apr 2026 11:29:04 +0100 Subject: [PATCH 4/4] Tidying up python and improving readability --- .../build/psyclone/psyclone_tools.py | 36 +++++++++++++------ 1 file changed, 25 insertions(+), 11 deletions(-) diff --git a/infrastructure/build/psyclone/psyclone_tools.py b/infrastructure/build/psyclone/psyclone_tools.py index 7c671cbb5..829733e90 100644 --- a/infrastructure/build/psyclone/psyclone_tools.py +++ b/infrastructure/build/psyclone/psyclone_tools.py @@ -99,10 +99,12 @@ def colour_loops(psyir: FileContainer, enable_tiling=False): # Colour loops over cells unless they are on discontinuous # spaces or over DoFs for child in subroutine.children: - # Check if the profiling calipers have been added before the colouring. + # Check if the profiling calipers have been added before the + # colouring. if isinstance(child, ProfileNode): raise TransformationError( - "Must apply colour_loops BEFORE profile_loops function in optimisation script.") + "Must apply colour_loops BEFORE profile_loops function " + "in optimisation script.") if ( isinstance(child, Loop) and child.iteration_space.endswith("cell_column") @@ -112,7 +114,7 @@ def colour_loops(psyir: FileContainer, enable_tiling=False): ctrans.apply(child, options={"tiling": enable_tiling}) # ----------------------------------------------------------------------------- -def profile_loops(psyir: FileContainer,colours_only=True): +def profile_loops(psyir: FileContainer, colours_only=True): """ Applies timing calipers to kernels during the psyclone build. The default is to only profile coloured loops but colours_only can be set to False to @@ -137,17 +139,29 @@ def profile_loops(psyir: FileContainer,colours_only=True): if not loop.coded_kernels(): continue # Insert profiler calls before loop over colours - if (not colours_only and not loop.loop_type in leave_loops) or loop.loop_type == "colours": - # First check that the transformation is not being made inside an OMP region. - if loop.ancestor(OMPParallelDirective) or loop.ancestor(OMPParallelDoDirective) \ - or loop.ancestor(OMPDoDirective): + if ((loop.loop_type == "colours") or + (colours_only is False and loop.loop_type not in leave_loops)): + # First check that the transformation is not being made inside + # an OMP region. + if (loop.ancestor(OMPParallelDirective) + or loop.ancestor(OMPParallelDoDirective) + or loop.ancestor(OMPDoDirective)): raise TransformationError( - "Must apply profile_loops BEFORE openmp_parellelise_loops function in optimisation script.") - k_name = loop.ancestor(InvokeSchedule).coded_kernels()[count].name + "Must apply profile_loops BEFORE " + "openmp_parallelise_loops function in optimisation " + "script.") + # Constructing unique calliper name based on kernel name, + # invoke name and kernel count + k_object = loop.ancestor(InvokeSchedule).coded_kernels()[count] + k_name = k_object.name invoke_name = loop.ancestor(InvokeSchedule).invoke.name file_name = loop.ancestor(Container).name - options = {"region_name": (file_name,invoke_name + ":" + k_name + "_k" + str(count))} - profile_trans.apply(loop,options=options) + # Make region name + region_name = invoke_name + ":" + k_name + "_k" + str(count) + options = {"region_name": (file_name, region_name)} + profile_trans.apply(loop, options=options) + # Count here is to distinguish kernels of the same name + # in the same invoke. count += 1 # -----------------------------------------------------------------------------