Skip to content

mosrs ticket 1048 - adding openmp directives to chunking of ukca full domain passing #172

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.

Sign up for GitHub

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

Open
RobWatersMet wants to merge 4 commits into
base: main
Choose a base branch
from
Open

mosrs ticket 1048 - adding openmp directives to chunking of ukca full domain passing #172

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
a84f833
moved changes over from mosrs ticket 1048, these were related to addi…
RobWatersMet Jan 20, 2026
2b24e38
adding myself to contributor list
RobWatersMet Jan 21, 2026
77a1f49
typo
RobWatersMet Jan 21, 2026
d21e513
Merge branch 'main' into ngarch_ukca_openmp_chunking_ticket_1048
RobWatersMet Jan 22, 2026
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
29 changes: 15 additions & 14 deletions CONTRIBUTORS.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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 |
| 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 |
12 changes: 6 additions & 6 deletions applications/lfric_atm/build/psyclone_transmute_file_list.mk
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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 =
Expand Down
136 changes: 127 additions & 9 deletions ...-sandbox/transmute/science/ukca/src/science/core/chemistry/ukca_chemistry_ctl_full_mod.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
# The file LICENCE, distributed with this code, contains details of the terms
# under which the code may be used.
##############################################################################

# Summary
# =======
#
Expand Down Expand Up @@ -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
# =======
#
Expand Down Expand Up @@ -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__)

Expand All @@ -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")
Expand All @@ -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:
Expand Down Expand Up @@ -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
Loading