From f54445ffc91383e1f9a4ede7def7250c727111f3 Mon Sep 17 00:00:00 2001 From: Matthew Naylor Date: Mon, 6 Jul 2026 10:57:22 +0100 Subject: [PATCH 1/5] Update PSyclone scripts for UKCA chemistry 3D chunking UKCA full-domain mode chemistry chunking is now implemented in the UKCA code base, not as an lfric_apps PSyclone transformer. The PSyclone transformer now only adds OpenMP directives to the already-chunked code. --- .../chemistry/ukca_chemistry_ctl_full_mod.py | 508 +++--------------- .../core/top_level/ukca_main1-ukca_main1.py | 108 ---- .../chemistry/ukca_chemistry_ctl_full_mod.py | 504 ++--------------- .../build/psyclone_transmute_file_list.mk | 3 +- 4 files changed, 121 insertions(+), 1002 deletions(-) delete mode 100644 applications/lfric_atm/optimisation/meto-ex1a/transmute/science/ukca/src/control/core/top_level/ukca_main1-ukca_main1.py 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 fb35cc63e..f8175cbe4 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,466 +3,80 @@ # The file LICENCE, distributed with this code, contains details of the terms # under which the code may be used. ############################################################################## -# Summary -# ======= -# -# This transformation introduces a chunking loop around the call to the ASAD -# solver in `ukca_chemistry_ctl_full_mod.F90`. A single call to the ASAD -# solver is replaced with multiple calls, each of which operates on a "chunk" -# of the full domain. The chunk size is taken at compile time from the -# environment variable UKCA_FULL_CHUNK_SIZE. If this variable is not set then -# the source code is passed through unmodified. -# -# Chunking is mainly achieved by slicing the arguments to the ASAD solver. -# However, the ASAD solver is dependent not only on its arguments but -# also on global ASAD arrays, several of which are read and written directly by -# UKCA full-domain mode. When chunking is enabled, any ASAD arrays that were -# originally full-domain sized become chunk sized. To resolve this size change, -# the transformer needs to be told about all of these ASAD arrays via -# the following parameter. -# -# * asad_vars: a dict mapping the ASAD arrays (originally -# full-domain sized), which are accessed by UKCA -# full-domain mode, to their associated ranks -# -# Unfortunately, this parameter cannot be inferred automatically because the -# ASAD arrays are dynamically allocated (and hence we don't know which ones -# were originally full-domain sized at compile time). Any changes to UKCA -# full-domain mode must therefore ensure that asad_vars is updated, if -# necessary. -# -# For each variable in asad_vars, the transformer introduces a -# full-domain-sized counterpart. Outside the chunking loop, it renames each -# access of a narrow (chunk sized) ASAD array into an access of its wide -# (full-domain sized) counterpart. Inside the chunking loop, slices of the -# newly introduced wide arrays are copied into narrow ASAD arrays, then the -# ASAD solver is called, and then narrow ASAD arrays are copied back into -# slices of their wide counterparts. -# -# In addition to asad_vars, the transformer uses the following parameters. -# -# * 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 parallel 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 -# ======= -# -# Given the program -# -# subroutine main() -# integer, parameter :: n = 256 -# integer :: arr1(n) -# integer :: arr2(n, 2) -# integer :: asad_arr1(32) -# arr1(:) = foo -# asad_arr1(:) = arr1(:) + 1 -# call asad_cdrive(arr1, arr2, n) -# arr1(:) = arr1(:) + asad_arr1(:) -# arr1(:) = arr1(:) + 1 -# end subroutine -# -# and the parameters -# -# fulldom_size_name = "n" -# asad_vars = {"asad_arr1" : 1} -# asad_call_name = "asad_cdrive" -# UKCA_FULL_CHUNK_SIZE = 32 -# -# the following program is produced. -# -# subroutine main() -# integer, parameter :: n = 256 -# integer :: arr1(n) -# integer :: arr2(n, 2) -# integer :: asad_arr1(32) -# integer :: full_asad_arr1(n) -# integer :: chunk_begin, chunk_end, chunk_size -# -# arr1(:) = foo -# full_asad_arr1(:) = arr1(:) + 1 -# do chunk_begin = 1, n, 32 -# chunk_end = min(n, chunk_begin+chunk_size-1) -# chunk_size = 1 + chunk_end - chunk_begin -# asad_arr1(1:chunk_size) = full_asad_arr1(chunk_begin:chunk_end) -# call asad_cdrive(arr1(chunk_begin:chunk_end), & -# arr2(chunk_begin:chunk_end, :), & -# chunk_size) -# full_asad_arr1(chunk_begin:chunk_end) = asad_arr1(1:chunk_size) -# end do -# arr1(:) = arr1(:) + full_asad_arr1(:) -# arr1(:) = arr1(:) + 1 -# end subroutine -# Imports -# ======= +# This transformation introduces OpenMP directives around loops inside +# UKCA full-domain mode. -import os import logging -from psyclone.psyir.nodes import ( - ArrayReference, - Assignment, - BinaryOperation, - Call, - IfBlock, - IntrinsicCall, - Literal, - Loop, - Reference, - Routine, - Schedule, - UnaryOperation, -) +import os + from psyclone.psyir.symbols import ( - CHARACTER_TYPE, - INTEGER_TYPE, - REAL_TYPE, ArrayType, - ContainerSymbol, DataSymbol, - ImportInterface, - RoutineSymbol, - Symbol, ) -from psyclone.psyir.transformations.reference2arrayrange_trans import ( - Reference2ArrayRangeTrans, +from psyclone.psyir.nodes import ( + Loop, + Reference, + Routine +) +from psyclone.transformations import ( + OMPLoopTrans, + TransformationError ) -from psyclone.transformations import OMPParallelLoopTrans, TransformationError from psyclone.version import __MAJOR__, __MICRO__, __MINOR__ -# Conditonal imports -# ================== - +# PSyclone version psy_version = (__MAJOR__, __MINOR__, __MICRO__) -# Transformation Parameters -# ========================= - -# Name of variable holding the full-domain size -fulldom_size_name = "tot_n_pnts" - -# ASAD arrays in use (and their ranks) -asad_vars = {"rk": 2, "sph2o": 1, "sphno3": 1, "tnd": 1, - "y": 2, "za": 1, "dpd": 2, "dpw": 2, - "prk": 2, "fpsc1": 1, "fpsc2": 1} - -# Name of the top-level ASAD call -asad_call_name = "asad_cdrive" - -# Name of the routine in which to apply the transformation -routine_name = "ukca_chemistry_ctl_full" - -# 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 match_loop(loop: Loop, var_name: str, stop_name: str) -> bool: + # Return true only if loop's variable is named var_name + # and loop's stop expression is a reference named stop_name. + return (loop.variable.name == var_name and + isinstance(loop.stop_expr, Reference) and + loop.stop_expr.name == stop_name) def trans(psyir): - desired_chunk_size = os.getenv("UKCA_FULL_CHUNK_SIZE") - if desired_chunk_size is None: - # Do nothing if the chunk size is not set - return - elif desired_chunk_size == "FULL_DOMAIN": - # Message to print (via umPrint) when chunking enabled - message_text = ("UKCA full-domain chunking enabled with " + - "a chunk size equal to the size of the full " + - "domain") - # We use None to represent the full-domain chunk size - desired_chunk_size = None - else: - # Message to print (via umPrint) when chunking enabled - message_text = ("UKCA full-domain chunking enabled with " + - "a chunk size of " + desired_chunk_size) + # All loops are dynamically scheduled + omp_trans = OMPLoopTrans(omp_directive="paralleldo", + omp_schedule="dynamic") - 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: - continue - - # Find variable holding the size of the full domain - try: - array_size_var = routine.symbol_table.lookup(fulldom_size_name) - except Exception: - continue - - # Find the call to the ASAD solver - asad_call = None - for call in routine.walk(Call): - if call.routine.name == asad_call_name: - if call.parent is routine: - asad_call = call - if asad_call is None: - continue - - # Find references to ASAD arrays before and after the call - # -------------------------------------------------------- - - refs_before = set() - refs_after = set() - for stmt in routine.children[:asad_call.position]: - for ref in stmt.walk(Reference): - if ref.name in asad_vars: - refs_before.add(ref.name) - for stmt in routine.children[asad_call.position+1:]: - for ref in stmt.walk(Reference): - if ref.name in asad_vars: - refs_after.add(ref.name) - - # Introduce variable to hold the desired chunk size - # ------------------------------------------------- - - desired_chunk_size_var = routine.symbol_table.find_or_create_tag( - "desired_chunk_size", - symbol_type=DataSymbol, - datatype=INTEGER_TYPE) - - if desired_chunk_size is None: - assign_desired_chunk_size = Assignment.create( - Reference(desired_chunk_size_var), - Reference(array_size_var)) - else: - assign_desired_chunk_size = Assignment.create( - Reference(desired_chunk_size_var), - Literal(str(desired_chunk_size), INTEGER_TYPE)) - - # Introduce full-domain array for each ASAD array - # ----------------------------------------------- - - full_vars = {} - for (var_name, var_rank) in asad_vars.items(): - # Create array bounds - bounds = [Reference(array_size_var)] - for i in range(2, var_rank+1): - bounds.append(IntrinsicCall.create( - IntrinsicCall.Intrinsic.SIZE, - [Reference(Symbol(var_name)), - ("dim", Literal(str(i), INTEGER_TYPE))])) - # Create variables - new_var = routine.symbol_table.find_or_create_tag( - "full_" + var_name, - symbol_type=DataSymbol, - datatype=ArrayType(REAL_TYPE, bounds)) - full_vars[var_name] = (bounds, new_var) - # Add initialiser - if var_name in refs_before: - initialiser = Assignment.create( - ArrayReference.create(new_var, [":" for b in bounds]), - Literal("0.0", REAL_TYPE)) - routine.addchild(initialiser, index=0) - - # Replace each use of ASAD array with full-domain counterpart - # ----------------------------------------------------------- - - for stmt in routine.children: - for ref in stmt.walk(Reference): - if ref.name in full_vars: - (_var_bounds, var_sym) = full_vars[ref.name] - ref.symbol = var_sym - - # Identify array references in call - # --------------------------------- - - for arg in asad_call.arguments: - for ref in arg.walk(Reference): - ref2arraytrans = Reference2ArrayRangeTrans() - if not isinstance(ref, ArrayReference): - if psy_version <= (3, 1, 0): - if ref.is_array: - ref2arraytrans.apply(ref) - else: - if ref.symbol.is_array: - ref2arraytrans.apply( - ref, allow_call_arguments=True) - - # Create a new "chunking" loop - # ---------------------------- - - # Create loop variables - chunk_begin_var = routine.symbol_table.find_or_create_tag( - "chunk_begin", - symbol_type=DataSymbol, - datatype=INTEGER_TYPE) - chunk_end_var = routine.symbol_table.find_or_create_tag( - "chunk_end", - symbol_type=DataSymbol, - datatype=INTEGER_TYPE) - chunk_size_var = routine.symbol_table.find_or_create_tag( - "chunk_size", - symbol_type=DataSymbol, - datatype=INTEGER_TYPE) - - # Create assignment for chunk_end - minop = IntrinsicCall.create( - IntrinsicCall.Intrinsic.MIN, - [Reference(array_size_var), - BinaryOperation.create( - BinaryOperation.Operator.ADD, - Reference(chunk_begin_var), - BinaryOperation.create( - BinaryOperation.Operator.SUB, - Reference(desired_chunk_size_var), - Literal("1", INTEGER_TYPE)))]) - assign_chunk_end = Assignment.create(Reference(chunk_end_var), minop) - - # Create assignment for chunk_size - chunk_size = BinaryOperation.create( - BinaryOperation.Operator.SUB, - BinaryOperation.create( - BinaryOperation.Operator.ADD, - Literal("1", INTEGER_TYPE), - Reference(chunk_end_var)), - Reference(chunk_begin_var)) - assign_chunk_size = Assignment.create(Reference(chunk_size_var), - chunk_size) - - # Create chunking loop - loop = Loop(variable=chunk_begin_var) - asad_call.replace_with(loop) - loop.children = [Literal("1", INTEGER_TYPE), - Reference(array_size_var), - Reference(desired_chunk_size_var), - Schedule(parent=loop, children=[asad_call])] - - # Copy full-size arrays into chunk-size arrays - for var_name in refs_before: - (bounds, full_var) = full_vars[var_name] - var_sym = DataSymbol( - var_name, datatype=ArrayType(REAL_TYPE, bounds)) - assign_full_var = Assignment.create( - ArrayReference.create(var_sym, [":" for b in bounds]), - ArrayReference.create(full_var, [":" for b in bounds])) - loop.loop_body.addchild(assign_full_var, index=asad_call.position) - - # Copy chunk-size arrays into full-size arrays - for var_name in refs_after: - (bounds, full_var) = full_vars[var_name] - var_sym = DataSymbol( - var_name, datatype=ArrayType(REAL_TYPE, bounds)) - assign_full_var = Assignment.create( - ArrayReference.create(full_var, [":" for b in bounds]), - ArrayReference.create(var_sym, [":" for b in bounds])) - loop.loop_body.addchild(assign_full_var) - - # Update references to fulldom_size_name - for ref in loop.loop_body.walk(Reference): - if ref.name == fulldom_size_name: - ref.replace_with(Reference(chunk_size_var)) - - # Update references to arrays - for ref in loop.loop_body.walk(ArrayReference): - if ref.name in asad_vars.keys(): - ref.indices[0].start = Literal("1", INTEGER_TYPE) - ref.indices[0].stop = Reference(chunk_size_var) - else: - ref.indices[0].start = Reference(chunk_begin_var) - ref.indices[0].stop = Reference(chunk_end_var) - - # Adding chunk assignments - loop.loop_body.addchild(assign_chunk_end, index=0) - loop.loop_body.addchild(assign_chunk_size, index=1) - - # Add print statement - # ------------------- - - print_call = Call() - print_call.addchild(Reference(RoutineSymbol("umPrint"))) - print_call.addchild(Literal(message_text, CHARACTER_TYPE)) - loop.parent.addchild(print_call, index=loop.position) - - # Assign desired chunk size - # ------------------------- - - 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_trans = 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_trans.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 + if routine.name == "ukca_chemistry_ctl_full": + for loop in routine.walk(Loop): + try: + # Parallelise the "DO l = 1, dim_ntp" loops + if match_loop(loop, "l", "dim_ntp"): + omp_trans.apply(loop) + + # Parallelise the "DO jspf = 1, jpcspf" loop + if match_loop(loop, "jspf", "jpcspf"): + omp_trans.apply(loop, force=True) + + # Parallelise the 3D chunking loop + if match_loop(loop, "zi", "model_levels"): + # Find all "chunk_" arrays (to be marked as private) + privates = set() + for sym in loop.get_all_accessed_symbols(): + if (sym.name.startswith("chunk_") and + isinstance(sym, DataSymbol) and + isinstance(sym.datatype, ArrayType)): + privates.add(sym) + + # Apply the transformation + parent, position = loop.parent, loop.position + omp_trans.apply(loop, force=True, collapse=3) + + # Mark explicitly private variables + if psy_version < (3, 3, 0): + loop.explicitly_private_symbols.update(privates) + else: + directive = parent.children[position] + directive.explicitly_private_symbols.update( + privates) + + except TransformationError as err: + err_msg = ("ukca_chemistry_ctl_full_mod.py: Error: " + "could not apply OMP transformation " + f"to loop '{loop.variable.name}': " + f"{err.message_text}") + raise TransformationError(err_msg) from err diff --git a/applications/lfric_atm/optimisation/meto-ex1a/transmute/science/ukca/src/control/core/top_level/ukca_main1-ukca_main1.py b/applications/lfric_atm/optimisation/meto-ex1a/transmute/science/ukca/src/control/core/top_level/ukca_main1-ukca_main1.py deleted file mode 100644 index 1d423f7ca..000000000 --- a/applications/lfric_atm/optimisation/meto-ex1a/transmute/science/ukca/src/control/core/top_level/ukca_main1-ukca_main1.py +++ /dev/null @@ -1,108 +0,0 @@ -############################################################################## -# (c) Crown copyright 2025 Met Office. All rights reserved. -# The file LICENCE, distributed with this code, contains details of the terms -# under which the code may be used. -############################################################################## - -# Summary -# ======= -# -# This transformation is used to adjust the number of points passed to the ASAD -# initialisation routine to account for full-domain chunking. The chunk size -# is taken at compile time from the environment variable UKCA_FULL_CHUNK_SIZE. -# If this variable is not set then the source code is passed through -# unmodified. -# -# Example -# ======= -# -# Give the program -# -# subroutine main() -# if (ukca_config%l_ukca_asad_full) then -# n_pnts = tot_n_pnts -# endif -# end subroutine -# -# and the parameters -# -# match_if = "ukca_config%l_ukca_asad_full" -# match_lhs = "n_pnts" -# match_rhs = "tot_n_pnts" -# UKCA_FULL_CHUNK_SIZE = 256 -# -# the following program is produced. -# -# subroutine main() -# if (ukca_config%l_ukca_asad_full) then -# n_pnts = 256 -# umPrint("Initialising ASAD solver with 256 points") -# end if -# end subroutine - -# Imports -# ======= - -import os - -from psyclone.psyir.nodes import ( - Assignment, Reference, Literal, IfBlock, Call) -from psyclone.psyir.symbols import ( - INTEGER_TYPE, RoutineSymbol, CHARACTER_TYPE) - -# Transformation Parameters -# ========================= - -# The if condition to match against -match_if = "ukca_config%l_ukca_asad_full" - -# The left-hand-side to match against -match_lhs = "n_pnts" - -# The right-hand-side to match against -match_rhs = "tot_n_pnts" - -# Transformation -# ============== - - -def trans(psyir): - chunk_size = os.getenv("UKCA_FULL_CHUNK_SIZE") - if chunk_size is None: - # Do nothing if the chunk size is not set - return - elif chunk_size == "FULL_DOMAIN": - # Message to print (via umPrint) when chunking enabled - message_text = "Initialising ASAD solver for full-domain" - # We use None to represent the full-domain chunk size - chunk_size = None - else: - # Message to print (via umPrint) when chunking enabled - message_text = ("Initialising ASAD solver with " + - chunk_size + - " points due to full-domain chunking") - - # Search and replace - found = None - for ifblock in psyir.walk(IfBlock): - if isinstance(ifblock.condition, Reference): - (sig, _inds) = ifblock.condition.get_signature_and_indices() - if str(sig) == match_if: - for assign in ifblock.walk(Assignment): - if ( - isinstance(assign.lhs, Reference) and - assign.lhs.name == match_lhs and - isinstance(assign.rhs, Reference) - ): - if assign.rhs.name == match_rhs: - if chunk_size is not None: - assign.rhs.replace_with( - Literal(str(chunk_size), INTEGER_TYPE)) - found = assign - - # Insert print call - if found: - print_call = Call() - print_call.addchild(Reference(RoutineSymbol("umPrint"))) - print_call.addchild(Literal(message_text, CHARACTER_TYPE)) - found.parent.addchild(print_call, index=found.position+1) 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 364c23851..f8175cbe4 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,466 +3,80 @@ # The file LICENCE, distributed with this code, contains details of the terms # under which the code may be used. ############################################################################## -# Summary -# ======= -# -# This transformation introduces a chunking loop around the call to the ASAD -# solver in `ukca_chemistry_ctl_full_mod.F90`. A single call to the ASAD -# solver is replaced with multiple calls, each of which operates on a "chunk" -# of the full domain. The chunk size is taken at compile time from the -# environment variable UKCA_FULL_CHUNK_SIZE. If this variable is not set then -# the source code is passed through unmodified. -# -# Chunking is mainly achieved by slicing the arguments to the ASAD solver. -# However, the ASAD solver is dependent not only on its arguments but -# also on global ASAD arrays, several of which are read and written directly by -# UKCA full-domain mode. When chunking is enabled, any ASAD arrays that were -# originally full-domain sized become chunk sized. To resolve this size change, -# the transformer needs to be told about all of these ASAD arrays via -# the following parameter. -# -# * asad_vars: a dict mapping the ASAD arrays (originally -# full-domain sized), which are accessed by UKCA -# full-domain mode, to their associated ranks -# -# Unfortunately, this parameter cannot be inferred automatically because the -# ASAD arrays are dynamically allocated (and hence we don't know which ones -# were originally full-domain sized at compile time). Any changes to UKCA -# full-domain mode must therefore ensure that asad_vars is updated, if -# necessary. -# -# For each variable in asad_vars, the transformer introduces a -# full-domain-sized counterpart. Outside the chunking loop, it renames each -# access of a narrow (chunk sized) ASAD array into an access of its wide -# (full-domain sized) counterpart. Inside the chunking loop, slices of the -# newly introduced wide arrays are copied into narrow ASAD arrays, then the -# ASAD solver is called, and then narrow ASAD arrays are copied back into -# slices of their wide counterparts. -# -# In addition to asad_vars, the transformer uses the following parameters. -# -# * 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 parallell 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 -# ======= -# -# Given the program -# -# subroutine main() -# integer, parameter :: n = 256 -# integer :: arr1(n) -# integer :: arr2(n, 2) -# integer :: asad_arr1(32) -# arr1(:) = foo -# asad_arr1(:) = arr1(:) + 1 -# call asad_cdrive(arr1, arr2, n) -# arr1(:) = arr1(:) + asad_arr1(:) -# arr1(:) = arr1(:) + 1 -# end subroutine -# -# and the parameters -# -# fulldom_size_name = "n" -# asad_vars = {"asad_arr1" : 1} -# asad_call_name = "asad_cdrive" -# UKCA_FULL_CHUNK_SIZE = 32 -# -# the following program is produced. -# -# subroutine main() -# integer, parameter :: n = 256 -# integer :: arr1(n) -# integer :: arr2(n, 2) -# integer :: asad_arr1(32) -# integer :: full_asad_arr1(n) -# integer :: chunk_begin, chunk_end, chunk_size -# -# arr1(:) = foo -# full_asad_arr1(:) = arr1(:) + 1 -# do chunk_begin = 1, n, 32 -# chunk_end = min(n, chunk_begin+chunk_size-1) -# chunk_size = 1 + chunk_end - chunk_begin -# asad_arr1(1:chunk_size) = full_asad_arr1(chunk_begin:chunk_end) -# call asad_cdrive(arr1(chunk_begin:chunk_end), & -# arr2(chunk_begin:chunk_end, :), & -# chunk_size) -# full_asad_arr1(chunk_begin:chunk_end) = asad_arr1(1:chunk_size) -# end do -# arr1(:) = arr1(:) + full_asad_arr1(:) -# arr1(:) = arr1(:) + 1 -# end subroutine -# Imports -# ======= +# This transformation introduces OpenMP directives around loops inside +# UKCA full-domain mode. import logging import os -from psyclone.psyir.nodes import ( - ArrayReference, - Assignment, - BinaryOperation, - Call, - IfBlock, - IntrinsicCall, - Literal, - Loop, - Reference, - Routine, - Schedule, - UnaryOperation, -) from psyclone.psyir.symbols import ( - CHARACTER_TYPE, - INTEGER_TYPE, - REAL_TYPE, ArrayType, - ContainerSymbol, DataSymbol, - ImportInterface, - RoutineSymbol, - Symbol, ) -from psyclone.psyir.transformations.reference2arrayrange_trans import ( - Reference2ArrayRangeTrans, +from psyclone.psyir.nodes import ( + Loop, + Reference, + Routine +) +from psyclone.transformations import ( + OMPLoopTrans, + TransformationError ) -from psyclone.transformations import OMPParallelLoopTrans, TransformationError from psyclone.version import __MAJOR__, __MICRO__, __MINOR__ -# Conditonal imports -# ================== - +# PSyclone version psy_version = (__MAJOR__, __MINOR__, __MICRO__) -# Transformation Parameters -# ========================= - -# Name of variable holding the full-domain size -fulldom_size_name = "tot_n_pnts" - -# ASAD arrays in use (and their ranks) -asad_vars = {"rk": 2, "sph2o": 1, "sphno3": 1, "tnd": 1, - "y": 2, "za": 1, "dpd": 2, "dpw": 2, - "prk": 2, "fpsc1": 1, "fpsc2": 1} - -# Name of the top-level ASAD call -asad_call_name = "asad_cdrive" - -# Name of the routine in which to apply the transformation -routine_name = "ukca_chemistry_ctl_full" - -# Source and name of the reallocation routine -asad_realloc_routine_loc = ("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 match_loop(loop: Loop, var_name: str, stop_name: str) -> bool: + # Return true only if loop's variable is named var_name + # and loop's stop expression is a reference named stop_name. + return (loop.variable.name == var_name and + isinstance(loop.stop_expr, Reference) and + loop.stop_expr.name == stop_name) def trans(psyir): - desired_chunk_size = os.getenv("UKCA_FULL_CHUNK_SIZE") - if desired_chunk_size is None: - return - elif desired_chunk_size == "FULL_DOMAIN": - # Message to print (via umPrint) when chunking enabled - message_text = ("UKCA full-domain chunking enabled with " + - "a chunk size equal to the size of the full " + - "domain") - # We use None to represent the full-domain chunk size - desired_chunk_size = None - else: - # Message to print (via umPrint) when chunking enabled - 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 + # All loops are dynamically scheduled + omp_trans = OMPLoopTrans(omp_directive="paralleldo", + omp_schedule="dynamic") - # Locate correct routine within which to apply the transformation for routine in psyir.walk(Routine): - if routine.name != routine_name: - continue - - # Find variable holding the size of the full domain - try: - array_size_var = routine.symbol_table.lookup(fulldom_size_name) - except Exception: - continue - - # Find the call to the ASAD solver - asad_call = None - for call in routine.walk(Call): - if call.routine.name == asad_call_name: - if call.parent is routine: - asad_call = call - if asad_call is None: - continue - - # Find references to ASAD arrays before and after the call - # -------------------------------------------------------- - - refs_before = set() - refs_after = set() - for stmt in routine.children[:asad_call.position]: - for ref in stmt.walk(Reference): - if ref.name in asad_vars: - refs_before.add(ref.name) - for stmt in routine.children[asad_call.position+1:]: - for ref in stmt.walk(Reference): - if ref.name in asad_vars: - refs_after.add(ref.name) - - # Introduce variable to hold the desired chunk size - # ------------------------------------------------- - - desired_chunk_size_var = routine.symbol_table.find_or_create_tag( - "desired_chunk_size", - symbol_type=DataSymbol, - datatype=INTEGER_TYPE) - - if desired_chunk_size is None: - assign_desired_chunk_size = Assignment.create( - Reference(desired_chunk_size_var), - Reference(array_size_var)) - else: - assign_desired_chunk_size = Assignment.create( - Reference(desired_chunk_size_var), - Literal(str(desired_chunk_size), INTEGER_TYPE)) - - # Introduce full-domain array for each ASAD array - # ----------------------------------------------- - - full_vars = {} - for (var_name, var_rank) in asad_vars.items(): - # Create array bounds - bounds = [Reference(array_size_var)] - for i in range(2, var_rank+1): - bounds.append(IntrinsicCall.create( - IntrinsicCall.Intrinsic.SIZE, - [Reference(Symbol(var_name)), - ("dim", Literal(str(i), INTEGER_TYPE))])) - # Create variables - new_var = routine.symbol_table.find_or_create_tag( - "full_" + var_name, - symbol_type=DataSymbol, - datatype=ArrayType(REAL_TYPE, bounds)) - full_vars[var_name] = (bounds, new_var) - # Add initialiser - if var_name in refs_before: - initialiser = Assignment.create( - ArrayReference.create(new_var, [":" for b in bounds]), - Literal("0.0", REAL_TYPE)) - routine.addchild(initialiser, index=0) - - # Replace each use of ASAD array with full-domain counterpart - # ----------------------------------------------------------- - - for stmt in routine.children: - for ref in stmt.walk(Reference): - if ref.name in full_vars: - (_var_bounds, var_sym) = full_vars[ref.name] - ref.symbol = var_sym - - # Identify array references in call - # --------------------------------- - - for arg in asad_call.arguments: - for ref in arg.walk(Reference): - ref2arraytrans = Reference2ArrayRangeTrans() - if not isinstance(ref, ArrayReference): - if psy_version <= (3, 1, 0): - if ref.is_array: - ref2arraytrans.apply(ref) - else: - if ref.symbol.is_array: - ref2arraytrans.apply( - ref, allow_call_arguments=True) - - # Create a new "chunking" loop - # ---------------------------- - - # Create loop variables - chunk_begin_var = routine.symbol_table.find_or_create_tag( - "chunk_begin", - symbol_type=DataSymbol, - datatype=INTEGER_TYPE) - chunk_end_var = routine.symbol_table.find_or_create_tag( - "chunk_end", - symbol_type=DataSymbol, - datatype=INTEGER_TYPE) - chunk_size_var = routine.symbol_table.find_or_create_tag( - "chunk_size", - symbol_type=DataSymbol, - datatype=INTEGER_TYPE) - - # Create assignment for chunk_end - minop = IntrinsicCall.create( - IntrinsicCall.Intrinsic.MIN, - [Reference(array_size_var), - BinaryOperation.create( - BinaryOperation.Operator.ADD, - Reference(chunk_begin_var), - BinaryOperation.create( - BinaryOperation.Operator.SUB, - Reference(desired_chunk_size_var), - Literal("1", INTEGER_TYPE)))]) - assign_chunk_end = Assignment.create(Reference(chunk_end_var), minop) - - # Create assignment for chunk_size - chunk_size = BinaryOperation.create( - BinaryOperation.Operator.SUB, - BinaryOperation.create( - BinaryOperation.Operator.ADD, - Literal("1", INTEGER_TYPE), - Reference(chunk_end_var)), - Reference(chunk_begin_var)) - assign_chunk_size = Assignment.create(Reference(chunk_size_var), - chunk_size) - - # Create chunking loop - loop = Loop(variable=chunk_begin_var) - asad_call.replace_with(loop) - loop.children = [Literal("1", INTEGER_TYPE), - Reference(array_size_var), - Reference(desired_chunk_size_var), - Schedule(parent=loop, children=[asad_call])] - - # Copy full-size arrays into chunk-size arrays - for var_name in refs_before: - (bounds, full_var) = full_vars[var_name] - var_sym = DataSymbol( - var_name, datatype=ArrayType(REAL_TYPE, bounds)) - assign_full_var = Assignment.create( - ArrayReference.create(var_sym, [":" for b in bounds]), - ArrayReference.create(full_var, [":" for b in bounds])) - loop.loop_body.addchild(assign_full_var, index=asad_call.position) - - # Copy chunk-size arrays into full-size arrays - for var_name in refs_after: - (bounds, full_var) = full_vars[var_name] - var_sym = DataSymbol( - var_name, datatype=ArrayType(REAL_TYPE, bounds)) - assign_full_var = Assignment.create( - ArrayReference.create(full_var, [":" for b in bounds]), - ArrayReference.create(var_sym, [":" for b in bounds])) - loop.loop_body.addchild(assign_full_var) - - # Update references to fulldom_size_name - for ref in loop.loop_body.walk(Reference): - if ref.name == fulldom_size_name: - ref.replace_with(Reference(chunk_size_var)) - - # Update references to arrays - for ref in loop.loop_body.walk(ArrayReference): - if ref.name in asad_vars.keys(): - ref.indices[0].start = Literal("1", INTEGER_TYPE) - ref.indices[0].stop = Reference(chunk_size_var) - else: - ref.indices[0].start = Reference(chunk_begin_var) - ref.indices[0].stop = Reference(chunk_end_var) - - # Adding chunk assignments - loop.loop_body.addchild(assign_chunk_end, index=0) - loop.loop_body.addchild(assign_chunk_size, index=1) - - # Add print statement - # ------------------- - - print_call = Call() - print_call.addchild(Reference(RoutineSymbol("umPrint"))) - print_call.addchild(Literal(message_text, CHARACTER_TYPE)) - loop.parent.addchild(print_call, index=loop.position) - - # Assign desired chunk size - # ------------------------- - - loop.parent.addchild(assign_desired_chunk_size, index=loop.position) - - if use_omp: - # Add import for ASAD reallocation routine - # ---------------------------------------- - sym_tab = routine.symbol_table - asad_realloc_mod_sym = sym_tab.find_or_create( - asad_realloc_routine_loc[0], - symbol_type=ContainerSymbol) - - asad_realloc_routine =sym_tab.find_or_create( - asad_realloc_routine_loc[1], - symbol_type=RoutineSymbol, - interface=ImportInterface(asad_realloc_mod_sym)) - - - # Add Reallocation Call to within Loop - # ------------------------------------ - realloc_call = Call.create( - asad_realloc_routine, - [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_trans = 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_trans.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 + if routine.name == "ukca_chemistry_ctl_full": + for loop in routine.walk(Loop): + try: + # Parallelise the "DO l = 1, dim_ntp" loops + if match_loop(loop, "l", "dim_ntp"): + omp_trans.apply(loop) + + # Parallelise the "DO jspf = 1, jpcspf" loop + if match_loop(loop, "jspf", "jpcspf"): + omp_trans.apply(loop, force=True) + + # Parallelise the 3D chunking loop + if match_loop(loop, "zi", "model_levels"): + # Find all "chunk_" arrays (to be marked as private) + privates = set() + for sym in loop.get_all_accessed_symbols(): + if (sym.name.startswith("chunk_") and + isinstance(sym, DataSymbol) and + isinstance(sym.datatype, ArrayType)): + privates.add(sym) + + # Apply the transformation + parent, position = loop.parent, loop.position + omp_trans.apply(loop, force=True, collapse=3) + + # Mark explicitly private variables + if psy_version < (3, 3, 0): + loop.explicitly_private_symbols.update(privates) + else: + directive = parent.children[position] + directive.explicitly_private_symbols.update( + privates) + + except TransformationError as err: + err_msg = ("ukca_chemistry_ctl_full_mod.py: Error: " + "could not apply OMP transformation " + f"to loop '{loop.variable.name}': " + f"{err.message_text}") + raise TransformationError(err_msg) from err diff --git a/applications/ngarch/build/psyclone_transmute_file_list.mk b/applications/ngarch/build/psyclone_transmute_file_list.mk index 51e3b2394..a1faa5ee8 100644 --- a/applications/ngarch/build/psyclone_transmute_file_list.mk +++ b/applications/ngarch/build/psyclone_transmute_file_list.mk @@ -40,8 +40,7 @@ export PSYCLONE_PHYSICS_FILES = \ sw_rad_tile_kernel_mod \ tr_mix \ ukca_aero_ctl \ - ukca_chemistry_ctl_full_mod \ - ukca_main1-ukca_main1 + ukca_chemistry_ctl_full_mod ##### TRANSMUTE_INCLUDE_METHOD specify_include ##### From c1bbf63d78b714ca694748f17a3cac83a4d322b9 Mon Sep 17 00:00:00 2001 From: Matthew Naylor Date: Mon, 6 Jul 2026 11:06:46 +0100 Subject: [PATCH 2/5] Add myself to contributors file --- CONTRIBUTORS.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index fdba723e2..2e54515a0 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -52,3 +52,4 @@ | marcstring | Marc Stringer | NCAS, Reading University | 2026-05-06 | | cameronbateman-mo | Cameron Bateman | Met Office | 2026-05-28 | | davelee2804 | David Lee | Bureau of Meteorology, Australia | 2026-06-02 | +| mn416 | Matthew Naylor | University of Cambridge | 2026-07-06 | From b599e01d7a3a4d771884f4cd1c24f6309c8326a0 Mon Sep 17 00:00:00 2001 From: Matthew Naylor Date: Mon, 6 Jul 2026 12:50:39 +0100 Subject: [PATCH 3/5] Remove PSyclone optimisations for `ukca_main1` --- .../build/psyclone_transmute_file_list.mk | 1 - .../core/top_level/ukca_main1-ukca_main1.py | 108 ------------------ .../core/top_level/ukca_main1-ukca_main1.py | 1 - .../build/psyclone_transmute_file_list.mk | 1 - 4 files changed, 111 deletions(-) delete mode 100644 applications/lfric_atm/optimisation/azngarch-sandbox/transmute/science/ukca/src/control/core/top_level/ukca_main1-ukca_main1.py delete mode 120000 applications/lfric_atm/optimisation/nci-gadi/transmute/science/ukca/src/control/core/top_level/ukca_main1-ukca_main1.py diff --git a/applications/lfric_atm/build/psyclone_transmute_file_list.mk b/applications/lfric_atm/build/psyclone_transmute_file_list.mk index 90e2556c0..bbe9d178e 100644 --- a/applications/lfric_atm/build/psyclone_transmute_file_list.mk +++ b/applications/lfric_atm/build/psyclone_transmute_file_list.mk @@ -40,7 +40,6 @@ export PSYCLONE_PHYSICS_FILES = \ 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 \ diff --git a/applications/lfric_atm/optimisation/azngarch-sandbox/transmute/science/ukca/src/control/core/top_level/ukca_main1-ukca_main1.py b/applications/lfric_atm/optimisation/azngarch-sandbox/transmute/science/ukca/src/control/core/top_level/ukca_main1-ukca_main1.py deleted file mode 100644 index 1d423f7ca..000000000 --- a/applications/lfric_atm/optimisation/azngarch-sandbox/transmute/science/ukca/src/control/core/top_level/ukca_main1-ukca_main1.py +++ /dev/null @@ -1,108 +0,0 @@ -############################################################################## -# (c) Crown copyright 2025 Met Office. All rights reserved. -# The file LICENCE, distributed with this code, contains details of the terms -# under which the code may be used. -############################################################################## - -# Summary -# ======= -# -# This transformation is used to adjust the number of points passed to the ASAD -# initialisation routine to account for full-domain chunking. The chunk size -# is taken at compile time from the environment variable UKCA_FULL_CHUNK_SIZE. -# If this variable is not set then the source code is passed through -# unmodified. -# -# Example -# ======= -# -# Give the program -# -# subroutine main() -# if (ukca_config%l_ukca_asad_full) then -# n_pnts = tot_n_pnts -# endif -# end subroutine -# -# and the parameters -# -# match_if = "ukca_config%l_ukca_asad_full" -# match_lhs = "n_pnts" -# match_rhs = "tot_n_pnts" -# UKCA_FULL_CHUNK_SIZE = 256 -# -# the following program is produced. -# -# subroutine main() -# if (ukca_config%l_ukca_asad_full) then -# n_pnts = 256 -# umPrint("Initialising ASAD solver with 256 points") -# end if -# end subroutine - -# Imports -# ======= - -import os - -from psyclone.psyir.nodes import ( - Assignment, Reference, Literal, IfBlock, Call) -from psyclone.psyir.symbols import ( - INTEGER_TYPE, RoutineSymbol, CHARACTER_TYPE) - -# Transformation Parameters -# ========================= - -# The if condition to match against -match_if = "ukca_config%l_ukca_asad_full" - -# The left-hand-side to match against -match_lhs = "n_pnts" - -# The right-hand-side to match against -match_rhs = "tot_n_pnts" - -# Transformation -# ============== - - -def trans(psyir): - chunk_size = os.getenv("UKCA_FULL_CHUNK_SIZE") - if chunk_size is None: - # Do nothing if the chunk size is not set - return - elif chunk_size == "FULL_DOMAIN": - # Message to print (via umPrint) when chunking enabled - message_text = "Initialising ASAD solver for full-domain" - # We use None to represent the full-domain chunk size - chunk_size = None - else: - # Message to print (via umPrint) when chunking enabled - message_text = ("Initialising ASAD solver with " + - chunk_size + - " points due to full-domain chunking") - - # Search and replace - found = None - for ifblock in psyir.walk(IfBlock): - if isinstance(ifblock.condition, Reference): - (sig, _inds) = ifblock.condition.get_signature_and_indices() - if str(sig) == match_if: - for assign in ifblock.walk(Assignment): - if ( - isinstance(assign.lhs, Reference) and - assign.lhs.name == match_lhs and - isinstance(assign.rhs, Reference) - ): - if assign.rhs.name == match_rhs: - if chunk_size is not None: - assign.rhs.replace_with( - Literal(str(chunk_size), INTEGER_TYPE)) - found = assign - - # Insert print call - if found: - print_call = Call() - print_call.addchild(Reference(RoutineSymbol("umPrint"))) - print_call.addchild(Literal(message_text, CHARACTER_TYPE)) - found.parent.addchild(print_call, index=found.position+1) diff --git a/applications/lfric_atm/optimisation/nci-gadi/transmute/science/ukca/src/control/core/top_level/ukca_main1-ukca_main1.py b/applications/lfric_atm/optimisation/nci-gadi/transmute/science/ukca/src/control/core/top_level/ukca_main1-ukca_main1.py deleted file mode 120000 index 1e508fee4..000000000 --- a/applications/lfric_atm/optimisation/nci-gadi/transmute/science/ukca/src/control/core/top_level/ukca_main1-ukca_main1.py +++ /dev/null @@ -1 +0,0 @@ -../../../../../../../../meto-ex1a/transmute/science/ukca/src/control/core/top_level/ukca_main1-ukca_main1.py \ No newline at end of file diff --git a/applications/linear_model/build/psyclone_transmute_file_list.mk b/applications/linear_model/build/psyclone_transmute_file_list.mk index 90e2556c0..bbe9d178e 100644 --- a/applications/linear_model/build/psyclone_transmute_file_list.mk +++ b/applications/linear_model/build/psyclone_transmute_file_list.mk @@ -40,7 +40,6 @@ export PSYCLONE_PHYSICS_FILES = \ 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 \ From 76061d86e3b753e1338d9e511162c11075f21543 Mon Sep 17 00:00:00 2001 From: Matthew Naylor Date: Mon, 6 Jul 2026 13:47:52 +0100 Subject: [PATCH 4/5] Apply suggestions from flake8 --- .../core/chemistry/ukca_chemistry_ctl_full_mod.py | 13 ++++++------- .../core/chemistry/ukca_chemistry_ctl_full_mod.py | 13 ++++++------- 2 files changed, 12 insertions(+), 14 deletions(-) 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 f8175cbe4..9c92f4d7d 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 @@ -7,9 +7,6 @@ # This transformation introduces OpenMP directives around loops inside # UKCA full-domain mode. -import logging -import os - from psyclone.psyir.symbols import ( ArrayType, DataSymbol, @@ -28,12 +25,14 @@ # PSyclone version psy_version = (__MAJOR__, __MINOR__, __MICRO__) + def match_loop(loop: Loop, var_name: str, stop_name: str) -> bool: # Return true only if loop's variable is named var_name # and loop's stop expression is a reference named stop_name. return (loop.variable.name == var_name and - isinstance(loop.stop_expr, Reference) and - loop.stop_expr.name == stop_name) + isinstance(loop.stop_expr, Reference) and + loop.stop_expr.name == stop_name) + def trans(psyir): # All loops are dynamically scheduled @@ -58,8 +57,8 @@ def trans(psyir): privates = set() for sym in loop.get_all_accessed_symbols(): if (sym.name.startswith("chunk_") and - isinstance(sym, DataSymbol) and - isinstance(sym.datatype, ArrayType)): + isinstance(sym, DataSymbol) and + isinstance(sym.datatype, ArrayType)): privates.add(sym) # Apply the transformation 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 f8175cbe4..9c92f4d7d 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 @@ -7,9 +7,6 @@ # This transformation introduces OpenMP directives around loops inside # UKCA full-domain mode. -import logging -import os - from psyclone.psyir.symbols import ( ArrayType, DataSymbol, @@ -28,12 +25,14 @@ # PSyclone version psy_version = (__MAJOR__, __MINOR__, __MICRO__) + def match_loop(loop: Loop, var_name: str, stop_name: str) -> bool: # Return true only if loop's variable is named var_name # and loop's stop expression is a reference named stop_name. return (loop.variable.name == var_name and - isinstance(loop.stop_expr, Reference) and - loop.stop_expr.name == stop_name) + isinstance(loop.stop_expr, Reference) and + loop.stop_expr.name == stop_name) + def trans(psyir): # All loops are dynamically scheduled @@ -58,8 +57,8 @@ def trans(psyir): privates = set() for sym in loop.get_all_accessed_symbols(): if (sym.name.startswith("chunk_") and - isinstance(sym, DataSymbol) and - isinstance(sym.datatype, ArrayType)): + isinstance(sym, DataSymbol) and + isinstance(sym.datatype, ArrayType)): privates.add(sym) # Apply the transformation From 4922e755d17f1004bf68c78b0010687d7d63ee7f Mon Sep 17 00:00:00 2001 From: Rob Waters Date: Tue, 7 Jul 2026 11:33:04 +0100 Subject: [PATCH 5/5] added ukca_chem_full_chunk_size to lfric ukca namelist --- .../rose-meta/um-chemistry/HEAD/rose-meta.conf | 13 +++++++++++++ .../source/support/um_ukca_init_mod.f90 | 4 +++- 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/interfaces/physics_schemes_interface/rose-meta/um-chemistry/HEAD/rose-meta.conf b/interfaces/physics_schemes_interface/rose-meta/um-chemistry/HEAD/rose-meta.conf index 22d7661b7..b8796221d 100644 --- a/interfaces/physics_schemes_interface/rose-meta/um-chemistry/HEAD/rose-meta.conf +++ b/interfaces/physics_schemes_interface/rose-meta/um-chemistry/HEAD/rose-meta.conf @@ -210,6 +210,19 @@ help=Use (True) to select passing full (MPI/ local) domain to the ns=namelist/Science/UM Chemistry sort-key=Panel-A02a type=logical +trigger=namelist:chemistry=ukca_chem_full_chunk_size: -1,-1,-1; + +[namelist:chemistry=ukca_chem_full_chunk_size] +type=integer +length=3 +range=-1: +description=Chunk size for full-domain ASAD chemistry solver +help=When l_ukca_asad_full=.true., ASAD processes the model + =domain in chunks of this size in (X, Y, nlev). Note Y is always 1 + =in LFRic. Use -1 for any dimension to use the full size of that dimension. +ns=namelist/Science/UM Chemistry +compulsory=false +sort-key=Panel-A02b [namelist:chemistry=l_ukca_linox_scaling] compulsory=true diff --git a/interfaces/physics_schemes_interface/source/support/um_ukca_init_mod.f90 b/interfaces/physics_schemes_interface/source/support/um_ukca_init_mod.f90 index f233d3085..07947e7eb 100644 --- a/interfaces/physics_schemes_interface/source/support/um_ukca_init_mod.f90 +++ b/interfaces/physics_schemes_interface/source/support/um_ukca_init_mod.f90 @@ -39,7 +39,8 @@ module um_ukca_init_mod photol_scheme_fastjx, & photol_scheme_prescribed, fastjx_mode, & fastjx_numwavel, fastjx_prescutoff, & - fjx_solcyc_type, fjx_solcyc_months + fjx_solcyc_type, fjx_solcyc_months, & + ukca_chem_full_chunk_size ! Other LFRic modules used use model_clock_mod, only: model_clock_type @@ -1863,6 +1864,7 @@ subroutine aerosol_ukca_dust_only_init( row_length, rows, model_levels, & ! i_mode_nzts=15, & ukca_mode_seg_size=i_ukca_mode_seg_size, & + ukca_chem_full_chunk_size=ukca_chem_full_chunk_size, & i_mode_setup=6, & i_mode_nucscav=i_mode_nucscav, & l_cv_rainout=.not.(l_ukca_plume_scav), &