Skip to content

Commit 3f088ce

Browse files
MetBenjaminWentsvadams
MetBenjaminWent
and
svadams
authored
Fix PR#377 Enable Transmute script for bl imp and exp kernels - over i (#534)
Co-authored-by: svadams <svadams@users.noreply.github.com>
1 parent 987533b commit 3f088ce

5 files changed

Lines changed: 176 additions & 61 deletions

File tree

‎applications/lfric_atm/build/psyclone_transmute_file_list.mk‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,9 @@
1717
export PSYCLONE_PHYSICS_FILES = \
1818
bl_lsp \
1919
bm_tau_kernel_mod \
20+
bl_exp_kernel_mod \
21+
bl_imp_kernel_mod \
22+
bl_imp2_kernel_mod \
2023
btq_int \
2124
conv_gr_kernel_mod \
2225
ex_flux_tq \

‎applications/lfric_atm/optimisation/meto-ex1a/transmute/global.py‎

Lines changed: 21 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,23 @@
1-
##############################################################################
2-
# Copyright (c) 2025, Met Office, on behalf of HMSO and Queen's Printer
3-
# For further details please refer to the file LICENCE.original which you
4-
# should have received as part of this distribution.
5-
##############################################################################
1+
# -----------------------------------------------------------------------------
2+
# (C) Crown copyright Met Office. All rights reserved.
3+
# The file LICENCE, distributed with this code, contains details of the terms
4+
# under which the code may be used.
5+
# -----------------------------------------------------------------------------
66
'''
77
A global script to add OpenMP to loops present in the file provided.
88
This script imports a SCRIPT_OPTIONS_DICT which can be used to override
99
small aspects of this script per file it is applied to.
1010
Overrides currently include:
11-
* Options list for transformations
12-
* safe pure calls for loops over calls which can be parallelised
11+
* ignore_dependencies_for
12+
* node_type_check
1313
'''
1414

1515
import logging
1616
from psyclone.transformations import (
1717
TransformationError)
1818
from psyclone.psyir.nodes import Loop
1919
from transmute_psytrans.transmute_functions import (
20-
OMP_PARALLEL_LOOP_DO_TRANS_STATIC,
21-
set_pure_subroutines,
20+
OMP_PARALLEL_LOOP_DO_TRANS_STATIC
2221
)
2322
from script_options import (
2423
SCRIPT_OPTIONS_DICT
@@ -32,32 +31,32 @@ def trans(psyir):
3231
to each loop.
3332
'''
3433

35-
# options list for transformation.
36-
options = {}
34+
fortran_file_name = str(psyir.root.name)
3735

38-
# Designate calls in regions as safe to parallelise over.
39-
safe_pure_calls = []
36+
node_type_check = True
37+
ignore_dependencies_for = []
4038

4139
fortran_file_name = str(psyir.root.name)
4240
# Check if file is in the script_options_dict
4341
# Copy out anything that's needed
4442
# options list and a pure calls override
4543
if fortran_file_name in SCRIPT_OPTIONS_DICT:
4644
file_overrides = SCRIPT_OPTIONS_DICT[fortran_file_name]
47-
if "options" in file_overrides.keys():
48-
options = file_overrides["options"]
49-
if "safe_pure_calls" in file_overrides.keys():
50-
safe_pure_calls = file_overrides["safe_pure_calls"]
51-
52-
# Set the pure calls if needed
53-
if safe_pure_calls:
54-
set_pure_subroutines(psyir, safe_pure_calls)
45+
if "ignore_dependencies_for" in file_overrides.keys():
46+
ignore_dependencies_for = file_overrides[
47+
"ignore_dependencies_for"]
48+
if "node_type_check" in file_overrides.keys():
49+
node_type_check = file_overrides[
50+
"node_type_check"]
5551

5652
# Work through each loop in the file and OMP PARALLEL DO
5753
for loop in psyir.walk(Loop):
5854
if not loop.ancestor(Loop):
5955
try:
60-
OMP_PARALLEL_LOOP_DO_TRANS_STATIC.apply(loop, options)
56+
OMP_PARALLEL_LOOP_DO_TRANS_STATIC.apply(
57+
loop,
58+
ignore_dependencies_for=ignore_dependencies_for,
59+
node_type_check=node_type_check)
6160

6261
except (TransformationError, IndexError) as err:
6362
logging.warning(

‎applications/lfric_atm/optimisation/meto-ex1a/transmute/kernel/local.py‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
# -----------------------------------------------------------------------------
2+
# (C) Crown copyright Met Office. All rights reserved.
3+
# The file LICENCE, distributed with this code, contains details of the terms
4+
# under which the code may be used.
5+
# -----------------------------------------------------------------------------
6+
'''
7+
A local.py script for all kernels, where instead of adding OMP across the
8+
outermost loop, it is placed around the i loop, or across the l loop.
9+
This script imports a SCRIPT_OPTIONS_DICT which can be used to override
10+
small aspects of this script per file it is applied to.
11+
Overrides currently include:
12+
* ignore_dependencies_for
13+
* node_type_check
14+
* safe_pure_calls
15+
'''
16+
17+
import logging
18+
from psyclone.transformations import (
19+
TransformationError)
20+
from psyclone.psyir.nodes import (
21+
Loop, Call,
22+
OMPParallelDoDirective,
23+
OMPParallelDirective,
24+
OMPDoDirective,)
25+
from transmute_psytrans.transmute_functions import (
26+
OMP_PARALLEL_LOOP_DO_TRANS_STATIC
27+
)
28+
from script_options import (
29+
SCRIPT_OPTIONS_DICT
30+
)
31+
32+
33+
def trans(psyir):
34+
'''
35+
PSyclone function call, run through psyir object,
36+
each schedule (or subroutine) and apply paralleldo transformations
37+
to each loop.
38+
'''
39+
40+
fortran_file_name = str(psyir.root.name)
41+
42+
node_type_check = True
43+
ignore_dependencies_for = []
44+
safe_pure_calls = []
45+
46+
if fortran_file_name in SCRIPT_OPTIONS_DICT:
47+
file_overrides = SCRIPT_OPTIONS_DICT[fortran_file_name]
48+
if "ignore_dependencies_for" in file_overrides.keys():
49+
ignore_dependencies_for = file_overrides[
50+
"ignore_dependencies_for"]
51+
if "node_type_check" in file_overrides.keys():
52+
node_type_check = file_overrides[
53+
"node_type_check"]
54+
if "safe_pure_calls" in file_overrides.keys():
55+
safe_pure_calls = file_overrides[
56+
"safe_pure_calls"]
57+
58+
# Set the calls to 'pure', given the provided override.
59+
# pure allows PSyclone to parallelise over them with OMP.
60+
if safe_pure_calls:
61+
for call in psyir.walk(Call):
62+
if call.routine.symbol.name in safe_pure_calls:
63+
call.routine.symbol.is_pure = True
64+
65+
# Work through each loop in the file and OMP PARALLEL DO
66+
for loop in psyir.walk(Loop):
67+
# If there is an OMP ancestor skip.
68+
if (
69+
loop.ancestor(OMPParallelDoDirective) is not None
70+
or loop.ancestor(OMPDoDirective) is not None
71+
or loop.ancestor(OMPParallelDirective) is not None
72+
):
73+
continue
74+
# Allow loops over 'i' and 'l' indexes to be parallelised.
75+
if loop.variable.name in ['i', 'l']:
76+
try:
77+
OMP_PARALLEL_LOOP_DO_TRANS_STATIC.apply(
78+
loop,
79+
ignore_dependencies_for=ignore_dependencies_for,
80+
node_type_check=node_type_check)
81+
except (TransformationError, IndexError) as err:
82+
logging.warning(
83+
"Could not transform because:\n %s", err)

‎applications/lfric_atm/optimisation/meto-ex1a/transmute/kernel/script_options.py‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
# -----------------------------------------------------------------------------
2+
# (C) Crown copyright Met Office. All rights reserved.
3+
# The file LICENCE, distributed with this code, contains details of the terms
4+
# under which the code may be used.
5+
# -----------------------------------------------------------------------------
6+
'''
7+
This file lifts optional overrides, where possible, into a localised,
8+
single location.
9+
The goal is to allow the creation of a simpler local.py which controls how OMP
10+
is added around loops.
11+
SCRIPT_OPTIONS_DICT is pulled into the local.py script, and overrides some
12+
local functionality.
13+
Each key represents a different kernel file which uses the local script.
14+
Each file key in turn contains a dictionary of overrides.
15+
'''
16+
17+
# Needs to be lifted and likely set by the build system longer term.
18+
# The filename passed to PSyclone, this is the pre-processed FTN source.
19+
FILE_EXTEN = ".xu90"
20+
21+
# Basic initialisation
22+
SCRIPT_OPTIONS_DICT = {}
23+
24+
# File keys
25+
SCRIPT_OPTIONS_DICT["bl_exp_kernel_mod"+str(FILE_EXTEN)] = {
26+
"ignore_dependencies_for": [
27+
"tnuc_nlcl", "dw_bl", "surf_interp", "bq_bl", "bt_bl",
28+
"dtrdz_tq_bl", "level_ent", "level_ent_dsc", "ent_we_lim",
29+
"ent_t_frac", "ent_zrzi", "ent_we_lim_dsc", "ent_t_frac_dsc",
30+
"ent_zrzi_dsc", "fd_taux", "fd_tauy", "gradrinr", "rhokm_bl",
31+
"rhokh_bl", "moist_flux_bl", "heat_flux_bl", "gradrinr", "lmix_bl",
32+
"ngstress_bl", "tke_bl", "bl_weight_1dbl", "zh_nonloc", "z_lcl",
33+
"inv_depth", "qcl_at_inv_top", "shallow_flag", "uw0_flux", "vw0_flux",
34+
"lcl_height", "parcel_top", "level_parcel_top", "wstar_2d", "thv_flux",
35+
"parcel_buoyancy", "qsat_at_lcl", "bl_type_ind", "visc_m_blend",
36+
"visc_h_blend", "zh_2d", "zhsc_2d", "ntml_2d", "cumulus_2d",
37+
"rh_crit", "mix_len_bm", "dsldzm", "wvar", "zht", "oblen",
38+
]
39+
}
40+
41+
SCRIPT_OPTIONS_DICT["bl_imp_kernel_mod"+str(FILE_EXTEN)] = {
42+
"ignore_dependencies_for": [
43+
"dqw_wth", "dtl_wth", "dqw_nt_wth", "dtl_nt_wth",
44+
"ct_ctq_wth", "qw_wth", "tl_wth", "dqw1_2d",
45+
"dtl1_2d", "ct_ctq1_2d",
46+
]
47+
}
48+
49+
SCRIPT_OPTIONS_DICT["bl_imp2_kernel_mod"+str(FILE_EXTEN)] = {
50+
"ignore_dependencies_for": [
51+
"fric_heating_blyr", "fric_heating_incv", "nblyr", "cca",
52+
"ccw", "z_lcl", "cf_area", "cf_bulk", "cf_ice", "cf_liq",
53+
"dtheta_bl", "m_v", "m_cl", "m_s", "fqw_star_w3", "ftl_star_w3",
54+
"heat_flux_bl", "moist_flux_bl", "rhokh_bl", "dqw_wth",
55+
"dtl_wth", "qw_wth", "tl_wth"
56+
]
57+
}
58+
59+
SCRIPT_OPTIONS_DICT["mphys_kernel_mod"+str(FILE_EXTEN)] = {
60+
"node_type_check": False,
61+
"ignore_dependencies_for": [
62+
"dml_wth", "dmv_wth", "dml_wth", "dtheta", "refl_1km", "dms_wth",
63+
"dmr_wth", "dmg_wth", "murk", "dbcf_wth", "dcff_wth", "dcfl_wth",
64+
"ls_rain_2d", "ls_snow_2d", "ls_graup_2d", "lsca_2d", "ls_rain_3d",
65+
"ls_snow_3d", "precfrac", "autoconv", "accretion", "rim_cry",
66+
"rim_agg", "refl_tot", "superc_liq_wth", "superc_rain_wth",
67+
"sfwater", "sfrain", "sfsnow",
68+
]
69+
}

‎applications/lfric_atm/optimisation/meto-ex1a/transmute/script_options.py‎

Lines changed: 0 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -16,42 +16,3 @@
1616

1717
# Basic initialisation, will be used by global script
1818
SCRIPT_OPTIONS_DICT = {}
19-
20-
# Kernels
21-
SCRIPT_OPTIONS_DICT["mphys_kernel_mod"+str(FILE_EXTEN)] = {
22-
23-
"options": {
24-
"node-type-check": False,
25-
"ignore_dependencies_for": [
26-
"dtheta", # First and Second i, k loop
27-
"dmv_wth", # First and Second i, k loop
28-
"dml_wth", # First and Second i, k loop
29-
"dms_wth", # First and Second i, k loop
30-
"dmr_wth", # Third i, k loop
31-
"dmg_wth", # Forth i, k loop
32-
"murk", # First k, i loop
33-
"dbcf_wth", # Fifth i, k loop
34-
"dcfl_wth", # Fifth i, k loop
35-
"dcff_wth", # Fifth i, k loop
36-
"ls_rain_2d", # Fifth i, k loop
37-
"ls_snow_2d", # Fifth i, k loop
38-
"ls_graup_2d", # Fifth i, k loop
39-
"lsca_2d", # Fifth i, k loop
40-
"ls_rain_3d", # Fifth i, k loop
41-
"ls_snow_3d", # Fifth i, k loop
42-
"precfrac", # Fifth i, k loop
43-
"refl_tot", # Fifth i, k loop
44-
"autoconv", # Fifth i, k loop
45-
"accretion", # Fifth i, k loop
46-
"rim_cry", # Fifth i, k loop
47-
"rim_agg", # Fifth i, k loop
48-
"refl_1km", # Fifth i, k loop
49-
"superc_liq_wth", # Sixth i, k loop
50-
"superc_rain_wth", # Seventh i, k loop
51-
"sfwater", # Eighth i, k loop
52-
"sfwater", # Second k, i loop
53-
"sfrain", # Third k, i loop
54-
"sfsnow", # Fourth k, i loop
55-
]
56-
}
57-
}

0 commit comments

Comments
 (0)