diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index 10e9ced23..1e43d6192 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -1,15 +1,16 @@ # Contributors -| GitHub user | Real Name | Affiliation | Date | -| --------------- | ----------------- | ----------- | ---------- | -| james-bruten-mo | James Bruten | Met Office | 2025-12-09 | -| jedbakerMO | Jed Baker | Met Office | 2025-12-29 | -| jennyhickson | Jenny Hickson | Met Office | 2025-12-10 | -| mike-hobson | Mike Hobson | Met Office | 2025-12-17 | -| mo-marqh | mark Hedley | Met Office | 2025-12-11 | -| yaswant | Yaswant Pradhan | Met Office | 2025-12-16 | -| oakleybrunt | Oakley Brunt | Met Office | 2025-12-19 | -| harry-shepherd | Harry Shepherd | Met Office | 2026-01-08 | -| DrTVockerodtMO | Terence Vockerodt | Met Office | 2026-01-08 | -| MetBenjaminWent | Benjamin Went | Met Office | 2026-01-15 | -| timgraham-Met | Tim Graham | Met Office | 2026-01-15 | -| mo-alistairp | Alistair Pirrie | Met Office | 2026-01-19 | \ No newline at end of file +| GitHub user | Real Name | Affiliation | Date | +| --------------- | ----------------- | ----------------------- | ---------- | +| james-bruten-mo | James Bruten | Met Office | 2025-12-09 | +| jedbakerMO | Jed Baker | Met Office | 2025-12-29 | +| jennyhickson | Jenny Hickson | Met Office | 2025-12-10 | +| mike-hobson | Mike Hobson | Met Office | 2025-12-17 | +| mo-marqh | mark Hedley | Met Office | 2025-12-11 | +| yaswant | Yaswant Pradhan | Met Office | 2025-12-16 | +| oakleybrunt | Oakley Brunt | Met Office | 2025-12-19 | +| harry-shepherd | Harry Shepherd | Met Office | 2026-01-08 | +| DrTVockerodtMO | Terence Vockerodt | Met Office | 2026-01-08 | +| MetBenjaminWent | Benjamin Went | Met Office | 2026-01-15 | +| timgraham-Met | Tim Graham | Met Office | 2026-01-15 | +| mo-alistairp | Alistair Pirrie | Met Office | 2026-01-19 | +| RobWatersMet | Rob Waters | University of Cambridge | 2026-01-21 | diff --git a/applications/lfric_atm/build/psyclone_transmute_file_list.mk b/applications/lfric_atm/build/psyclone_transmute_file_list.mk index 215126425..8f12df2d9 100644 --- a/applications/lfric_atm/build/psyclone_transmute_file_list.mk +++ b/applications/lfric_atm/build/psyclone_transmute_file_list.mk @@ -21,14 +21,14 @@ export PSYCLONE_PHYSICS_FILES = mphys_kernel_mod \ pc2_initiation_kernel_mod \ pc2_conv_coupling_kernel_mod \ sw_kernel_mod \ - ukca_aero_ctl \ - ukca_abdulrazzak_ghan \ + ukca_aero_ctl \ + ukca_abdulrazzak_ghan \ ukca_chemistry_ctl_full_mod \ ukca_main1-ukca_main1 \ sw_rad_tile_kernel_mod \ - jules_imp_kernel_mod \ - jules_exp_kernel_mod \ - jules_extra_kernel_mod + jules_imp_kernel_mod \ + jules_exp_kernel_mod \ + jules_extra_kernel_mod ##### TRANSMUTE_INCLUDE_METHOD specify_exclude ##### # For GPU, we may want to use more generic local.py transformation scripts and psyclone by directory. @@ -38,7 +38,7 @@ export PSYCLONE_PHYSICS_FILES = mphys_kernel_mod \ # These files will be filtered, and will NOT be run through PSyclone. # Directories to psyclone -export PSYCLONE_DIRECTORIES = science/ukca +export PSYCLONE_DIRECTORIES = # A general file exception list export PSYCLONE_PHYSICS_EXCEPTION = diff --git a/applications/lfric_atm/optimisation/azngarch-sandbox/transmute/science/ukca/src/science/core/chemistry/ukca_chemistry_ctl_full_mod.py b/applications/lfric_atm/optimisation/azngarch-sandbox/transmute/science/ukca/src/science/core/chemistry/ukca_chemistry_ctl_full_mod.py index ac0e4ac2a..bfb1e4e68 100644 --- a/applications/lfric_atm/optimisation/azngarch-sandbox/transmute/science/ukca/src/science/core/chemistry/ukca_chemistry_ctl_full_mod.py +++ b/applications/lfric_atm/optimisation/azngarch-sandbox/transmute/science/ukca/src/science/core/chemistry/ukca_chemistry_ctl_full_mod.py @@ -3,7 +3,6 @@ # The file LICENCE, distributed with this code, contains details of the terms # under which the code may be used. ############################################################################## - # Summary # ======= # @@ -45,6 +44,19 @@ # * fulldom_size_name: name of variable holding full-domain size # * asad_call_name: name of the top-level ASAD solver routine # +# OpenMP can also be added to the chunked loop by setting the environment +# variable UKCA_FULL_CHUNK_OMP to True. By default it will be turned on +# provided the chunk size is not equal to domain size, i.e. loop of length 1 +# +# OpenMP parallelism is then added to the chunking loop using an omp parllel do +# directive to allow for top level parallelism on the ASAD solver. Importantly +# a call to ukca_reallocate_asad_arrays which reallocates the THREADPRIVATE +# arrays. This is done within the parallel region to account for the potential +# for chunk_size to be different between iterations, i.e. smaller last +# iteration. Dynamic scheduling has been selected based upon the number of +# solver iterations varying between chunks depending on the complexity of the +# chemistry. +# # Example # ======= # @@ -98,16 +110,40 @@ # ======= import os - -from psyclone.version import (__MAJOR__, __MINOR__, __MICRO__) +import logging from psyclone.psyir.nodes import ( - ArrayReference, Assignment, BinaryOperation, Call, IntrinsicCall, - Literal, Loop, Routine, Reference, Schedule) + ArrayReference, + Assignment, + BinaryOperation, + Call, + IfBlock, + IntrinsicCall, + Literal, + Loop, + Reference, + Routine, + Schedule, + UnaryOperation, +) from psyclone.psyir.symbols import ( - INTEGER_TYPE, REAL_TYPE, CHARACTER_TYPE, Symbol, DataSymbol, ArrayType, - RoutineSymbol) + CHARACTER_TYPE, + INTEGER_TYPE, + REAL_TYPE, + ArrayType, + ContainerSymbol, + DataSymbol, + ImportInterface, + RoutineSymbol, + Symbol, +) from psyclone.psyir.transformations.reference2arrayrange_trans import ( - Reference2ArrayRangeTrans) + Reference2ArrayRangeTrans, +) +from psyclone.transformations import OMPParallelLoopTrans, TransformationError +from psyclone.version import __MAJOR__, __MICRO__, __MINOR__ + +# Conditonal imports +# ================== psy_version = (__MAJOR__, __MINOR__, __MICRO__) @@ -128,9 +164,22 @@ # Name of the routine in which to apply the transformation routine_name = "ukca_chemistry_ctl_full" -# Transformation +# Source and name of the reallocation routine +asad_realloc_routine = ("ukca_chemistry_ctl_col_mod", + "ukca_reallocate_asad_arrays") + + +# Utility # ============== +def get_bool_env(var_name: str, default: bool = False) -> bool: + val = os.getenv(var_name) + if val is None: + return default + return val.strip().lower() in ('1', 'true', 't', 'yes', 'y', 'on') + +# Transformation +# ============== def trans(psyir): desired_chunk_size = os.getenv("UKCA_FULL_CHUNK_SIZE") @@ -149,6 +198,12 @@ def trans(psyir): message_text = ("UKCA full-domain chunking enabled with " + "a chunk size of " + desired_chunk_size) + use_omp = get_bool_env("UKCA_FULL_CHUNK_OMP", True) + if desired_chunk_size is None and use_omp: + logging.WARNING( + "Turning off omp as chunk size is set to full domain size") + use_omp = False + # Locate correct routine within which to apply the transformation for routine in psyir.walk(Routine): if routine.name != routine_name: @@ -348,3 +403,66 @@ def trans(psyir): # ------------------------- loop.parent.addchild(assign_desired_chunk_size, index=loop.position) + + if use_omp: + # Add import for ASAD reallocation routine + # ---------------------------------------- + + asad_realloc_mod_sym = ContainerSymbol(asad_realloc_routine[0]) + routine.symbol_table.add(asad_realloc_mod_sym) + asad_realloc_routine_sym = Symbol(asad_realloc_routine[1]) + asad_realloc_routine_sym.interface = ImportInterface( + asad_realloc_mod_sym) + routine.symbol_table.add(asad_realloc_routine_sym) + + # Add Reallocation Call to within Loop + # ------------------------------------ + # Create reallocation call + realloc_call = Call() + realloc_call.addchild(Reference(RoutineSymbol( + asad_realloc_routine[1]))) + realloc_call.addchild(Reference(chunk_size_var)) + + # Create conditional reallocation call + realloc_block = IfBlock.create( + BinaryOperation.create( + BinaryOperation.Operator.OR, + UnaryOperation.create( + UnaryOperation.Operator.NOT, + IntrinsicCall.create( + IntrinsicCall.Intrinsic.ALLOCATED, + [Reference( + Symbol(next(iter(asad_vars.keys()))))])), + BinaryOperation.create( + BinaryOperation.Operator.NE, + Reference(chunk_size_var), + IntrinsicCall.create( + IntrinsicCall.Intrinsic.SIZE, + [Reference(Symbol(next(iter(asad_vars.keys())))), + ("dim", Literal("1", INTEGER_TYPE))]))), + [realloc_call]) + + loop.loop_body.addchild(realloc_block, index=2) + + # Added OMP transformation on desired loop + # ---------------------------------------- + + omp_tans = OMPParallelLoopTrans(omp_schedule="static") + opts = { + # some non-PURE subroutines called within this loop + "force": True, + # several WRITE statements used for diagnostics + "node-type-check": False, + } + + try: + omp_tans.apply( + loop, options=opts, + ) + + except TransformationError as err: + err_msg = ("ukca_chemistry_ctl_full_mod.py: Error: " + "could not apply OMP transformation " + f"to loop: {err.message_text}") + + raise TransformationError(err_msg) from err diff --git a/applications/lfric_atm/optimisation/meto-ex1a/transmute/science/ukca/src/science/core/chemistry/ukca_chemistry_ctl_full_mod.py b/applications/lfric_atm/optimisation/meto-ex1a/transmute/science/ukca/src/science/core/chemistry/ukca_chemistry_ctl_full_mod.py index ac0e4ac2a..bfb1e4e68 100644 --- a/applications/lfric_atm/optimisation/meto-ex1a/transmute/science/ukca/src/science/core/chemistry/ukca_chemistry_ctl_full_mod.py +++ b/applications/lfric_atm/optimisation/meto-ex1a/transmute/science/ukca/src/science/core/chemistry/ukca_chemistry_ctl_full_mod.py @@ -3,7 +3,6 @@ # The file LICENCE, distributed with this code, contains details of the terms # under which the code may be used. ############################################################################## - # Summary # ======= # @@ -45,6 +44,19 @@ # * fulldom_size_name: name of variable holding full-domain size # * asad_call_name: name of the top-level ASAD solver routine # +# OpenMP can also be added to the chunked loop by setting the environment +# variable UKCA_FULL_CHUNK_OMP to True. By default it will be turned on +# provided the chunk size is not equal to domain size, i.e. loop of length 1 +# +# OpenMP parallelism is then added to the chunking loop using an omp parllel do +# directive to allow for top level parallelism on the ASAD solver. Importantly +# a call to ukca_reallocate_asad_arrays which reallocates the THREADPRIVATE +# arrays. This is done within the parallel region to account for the potential +# for chunk_size to be different between iterations, i.e. smaller last +# iteration. Dynamic scheduling has been selected based upon the number of +# solver iterations varying between chunks depending on the complexity of the +# chemistry. +# # Example # ======= # @@ -98,16 +110,40 @@ # ======= import os - -from psyclone.version import (__MAJOR__, __MINOR__, __MICRO__) +import logging from psyclone.psyir.nodes import ( - ArrayReference, Assignment, BinaryOperation, Call, IntrinsicCall, - Literal, Loop, Routine, Reference, Schedule) + ArrayReference, + Assignment, + BinaryOperation, + Call, + IfBlock, + IntrinsicCall, + Literal, + Loop, + Reference, + Routine, + Schedule, + UnaryOperation, +) from psyclone.psyir.symbols import ( - INTEGER_TYPE, REAL_TYPE, CHARACTER_TYPE, Symbol, DataSymbol, ArrayType, - RoutineSymbol) + CHARACTER_TYPE, + INTEGER_TYPE, + REAL_TYPE, + ArrayType, + ContainerSymbol, + DataSymbol, + ImportInterface, + RoutineSymbol, + Symbol, +) from psyclone.psyir.transformations.reference2arrayrange_trans import ( - Reference2ArrayRangeTrans) + Reference2ArrayRangeTrans, +) +from psyclone.transformations import OMPParallelLoopTrans, TransformationError +from psyclone.version import __MAJOR__, __MICRO__, __MINOR__ + +# Conditonal imports +# ================== psy_version = (__MAJOR__, __MINOR__, __MICRO__) @@ -128,9 +164,22 @@ # Name of the routine in which to apply the transformation routine_name = "ukca_chemistry_ctl_full" -# Transformation +# Source and name of the reallocation routine +asad_realloc_routine = ("ukca_chemistry_ctl_col_mod", + "ukca_reallocate_asad_arrays") + + +# Utility # ============== +def get_bool_env(var_name: str, default: bool = False) -> bool: + val = os.getenv(var_name) + if val is None: + return default + return val.strip().lower() in ('1', 'true', 't', 'yes', 'y', 'on') + +# Transformation +# ============== def trans(psyir): desired_chunk_size = os.getenv("UKCA_FULL_CHUNK_SIZE") @@ -149,6 +198,12 @@ def trans(psyir): message_text = ("UKCA full-domain chunking enabled with " + "a chunk size of " + desired_chunk_size) + use_omp = get_bool_env("UKCA_FULL_CHUNK_OMP", True) + if desired_chunk_size is None and use_omp: + logging.WARNING( + "Turning off omp as chunk size is set to full domain size") + use_omp = False + # Locate correct routine within which to apply the transformation for routine in psyir.walk(Routine): if routine.name != routine_name: @@ -348,3 +403,66 @@ def trans(psyir): # ------------------------- loop.parent.addchild(assign_desired_chunk_size, index=loop.position) + + if use_omp: + # Add import for ASAD reallocation routine + # ---------------------------------------- + + asad_realloc_mod_sym = ContainerSymbol(asad_realloc_routine[0]) + routine.symbol_table.add(asad_realloc_mod_sym) + asad_realloc_routine_sym = Symbol(asad_realloc_routine[1]) + asad_realloc_routine_sym.interface = ImportInterface( + asad_realloc_mod_sym) + routine.symbol_table.add(asad_realloc_routine_sym) + + # Add Reallocation Call to within Loop + # ------------------------------------ + # Create reallocation call + realloc_call = Call() + realloc_call.addchild(Reference(RoutineSymbol( + asad_realloc_routine[1]))) + realloc_call.addchild(Reference(chunk_size_var)) + + # Create conditional reallocation call + realloc_block = IfBlock.create( + BinaryOperation.create( + BinaryOperation.Operator.OR, + UnaryOperation.create( + UnaryOperation.Operator.NOT, + IntrinsicCall.create( + IntrinsicCall.Intrinsic.ALLOCATED, + [Reference( + Symbol(next(iter(asad_vars.keys()))))])), + BinaryOperation.create( + BinaryOperation.Operator.NE, + Reference(chunk_size_var), + IntrinsicCall.create( + IntrinsicCall.Intrinsic.SIZE, + [Reference(Symbol(next(iter(asad_vars.keys())))), + ("dim", Literal("1", INTEGER_TYPE))]))), + [realloc_call]) + + loop.loop_body.addchild(realloc_block, index=2) + + # Added OMP transformation on desired loop + # ---------------------------------------- + + omp_tans = OMPParallelLoopTrans(omp_schedule="static") + opts = { + # some non-PURE subroutines called within this loop + "force": True, + # several WRITE statements used for diagnostics + "node-type-check": False, + } + + try: + omp_tans.apply( + loop, options=opts, + ) + + except TransformationError as err: + err_msg = ("ukca_chemistry_ctl_full_mod.py: Error: " + "could not apply OMP transformation " + f"to loop: {err.message_text}") + + raise TransformationError(err_msg) from err