A personal note for a Python/HPC programmer who has to optimise LFRic
Claude Code
The Fortran note ended with a hand-written PSy layer and the claim that PSyclone is “a tracing JIT that rewrites source ahead of time”. This note cashes that claim. It is a rebuild of PSyclone in dependency order, aimed at two jobs: reading the generated Fortran in an LFRic build tree without treating it as magic, and writing the optimisation scripts that produce it — which is what a performance ticket against lfric_apps actually consists of.
The order is deliberate. PSyclone is usually taught from its DSL, because that is the historical entry point; here the DSL comes fourth, after the intermediate representation it is built on. That is the right order for someone who already knows what an IR and a compiler pass are, and it is also the order in which the tool now factors: PSyIR is the substrate, and the LFRic DSL is one client of it.
The essence. PSyclone is a compiler whose output has to be maintainable Fortran. Every design decision follows from that one constraint. It explains why the IR is unusually high-level (a Loop node, not a basic-block CFG — you cannot print a CFG back out as readable do loops); why there is an escape hatch node for constructs it cannot model (CodeBlock, because refusing to parse a 3,000-line physics routine over one write statement would make the tool useless); why transformations refuse rather than silently miscompile (a wrong !$omp in a checked-in file is a bug with a two-year half-life); and why the whole apparatus is driven by a Python script you commit next to the source, instead of compiler flags — because the optimisation is source, and gets reviewed like source.1
Everything below is executed. There is no cell magic this time and no compiler in the loop for most of the page: PSyclone is a Python library, so the notebook imports it and drives it in-process. The small helper module (src/psyclone_tour/) exists for presentation only — show prints Fortran as a highlighted block, tree prints a PSyIR tree, diff prints the delta a transformation made, run shells out for the parts that only exist as a command line.
from importlib.metadata import version
from psyclone_tour import LFRIC_APPS, LFRIC_CORE, PSYCLONE_SRC, diff, fortran, psyir, run, show, tree
from psyclone_tour.paths import LFRIC_CONFIG, LFRIC_PSYCLONE_TOOLS, LFRIC_TRANSMUTE_TOOLS, check
check() # fails with a sentence, not a stack trace, if the submodules are absent
print("psyclone", version("psyclone"), "| fparser", version("fparser"))psyclone 3.3.1 | fparser 0.2.4Which version, and why it matters. PSyclone’s Python API is not stable across minor releases, and LFRic tracks it closely rather than pinning far behind. lfric_apps main currently builds with 3.3.1 (spack load [email protected] in the ESNZ cascade suite), having migrated in “Updates needed to support psyclone 3.3”; that migration moved OMPParallelTrans into psyclone.psyir.transformations and renamed every Dynamo0p3*Trans to LFRic*Trans. This environment therefore pins the 3.3 series, so that the real optimisation scripts quoted below import and run unmodified. If you read an older script — or an older ticket — expect the import lines to have moved.2
A source-to-source Fortran compiler, distributed as a Python package, that reads Fortran, builds a tree, lets you mutate the tree from a script, and writes Fortran back out. It generates no object code and owns no runtime. The compiler you already use runs afterwards, on its output.
Fortran in
│
▼
fparser2 ──────────► a faithful, low-level parse tree
│ (every comma and keyword, no semantics)
▼
PSyIR ──────────► a semantic tree: Loops, Assignments,
│ Symbols, and CodeBlocks where it gave up
▼
your trans(psyir) ────► mutation, in place, validated
│
▼
FortranWriter ───────► Fortran outFour stages, and the practical consequences of each are different, so it is worth being precise about who is responsible for what:
| stage | what it owns | how it fails you |
|---|---|---|
| fparser2 | syntax | a syntax error, or a construct it cannot parse at all |
| PSyIR frontend | meaning: symbols, types, control flow | silently produces a CodeBlock — parsed, but opaque |
| transformations | legality and the rewrite itself | raises TransformationError and changes nothing |
| FortranWriter | printing | reformats aggressively; drops comments unless told not to |
PSyclone is really two tools sharing an IR, and LFRic uses both. Keeping them apart is the single most useful thing to fix early, because almost every confusing error message is a symptom of being in one mode while thinking in the other.
PSyKAl mode (-api lfric) |
transformation mode (no -api) |
|
|---|---|---|
| input | a .x90 algorithm layer with call invoke(...) |
any Fortran file |
| PSyclone generates code | yes — the whole PSy layer, from metadata | no — it only rewrites what you gave it |
| what the script sees | the PSy layer PSyclone just invented | the routine you wrote |
| parallelism knows about | function spaces, dofs, halos, colouring | loops, arrays, dependences |
| in LFRic | lfric_core, gungho, anything .x90 |
UM-inherited physics, via “transmute” |
| the makefile | psyclone_psykal.mk |
psyclone_transmute.mk |
| in this note | §5 | §6 |
Both are driven the same way: psyclone -s script.py, where the script defines exactly one function, trans(psyir).
The Fortran note gave the JAX correspondence for the PSyKAl layers. Here is the one for the tool, which is a different and more exact analogy — PSyclone is much closer to Numba’s and JAX’s internals than to a preprocessor.
| you know | PSyclone | where it differs |
|---|---|---|
jaxpr / Numba IR |
PSyIR | must print back out as readable Fortran, so it stays statement-shaped: no SSA, no CFG, no basic blocks |
| tracing a function | FortranReader.psyir_from_source |
it is parsing, not tracing: no values, so no constant folding for free, and control flow survives intact |
a JAX transform (vmap, grad) |
a Transformation subclass |
mutates one tree in place rather than returning a new function; and it validates first, refusing when illegal |
jax.jit lowering to XLA |
FortranWriter |
the target is text a human will read in a code review |
| XLA fusion decisions | LoopFuseTrans and friends, applied by hand from your script |
nothing is automatic; PSyclone applies exactly the transformations your script names |
jax.grad |
PSyAD (§7) | source-to-source reverse mode, ahead of time |
numba.objmode |
CodeBlock |
the escape hatch for what the IR cannot model |
The last row of the “differs” column is the one to internalise. PSyclone has no optimiser. It has a library of legal rewrites and a script in which you choose them. There is no -O2. Ticket #813 is 124 lines of Python because someone had to decide, loop by loop, what to do — and the deciding, not the typing, is the work.
Start where everything starts: text in, tree out.
source = """
module diffuse_mod
use constants_mod, only : r_def, i_def
implicit none
contains
subroutine diffuse(field, kappa, n, nlayers)
integer(i_def), intent(in) :: n, nlayers
real(r_def), intent(in) :: kappa
real(r_def), intent(inout) :: field(n, nlayers)
integer(i_def) :: i, k
do k = 1, nlayers
do i = 2, n - 1
field(i,k) = field(i,k) + kappa * (field(i-1,k) - 2.0_r_def*field(i,k) + field(i+1,k))
end do
end do
end subroutine diffuse
end module diffuse_mod
"""
tree(psyir(source))FileContainer[]
Container[diffuse_mod]
Routine[name:'diffuse']
0: Loop[variable='k']
Literal[value:'1', Scalar<INTEGER, UNDEFINED>]
Reference[name:'nlayers']
Literal[value:'1', Scalar<INTEGER, UNDEFINED>]
Schedule[]
0: Loop[variable='i']
Literal[value:'2', Scalar<INTEGER, UNDEFINED>]
BinaryOperation[operator:'SUB']
Reference[name:'n']
Literal[value:'1', Scalar<INTEGER, UNDEFINED>]
Literal[value:'1', Scalar<INTEGER, UNDEFINED>]
Schedule[]
0: Assignment[]
ArrayReference[name:'field']
Reference[name:'i']
Reference[name:'k']
BinaryOperation[operator:'ADD']
ArrayReference[name:'field']
Reference[name:'i']
Reference[name:'k']
BinaryOperation[operator:'MUL']
Reference[name:'kappa']
BinaryOperation[operator:'ADD']
BinaryOperation[operator:'SUB']
ArrayReference[name:'field']
BinaryOperation[operator:'SUB']
Reference[name:'i']
Literal[value:'1', Scalar<INTEGER, UNDEFINED>]
Reference[name:'k']
BinaryOperation[operator:'MUL']
Literal[value:'2.0', Scalar<REAL, Reference[name:'r_def']>]
ArrayReference[name:'field']
Reference[name:'i']
Reference[name:'k']
ArrayReference[name:'field']
BinaryOperation[operator:'ADD']
Reference[name:'i']
Literal[value:'1', Scalar<INTEGER, UNDEFINED>]
Reference[name:'k']Read that tree structurally, because four things about its shape are load-bearing.
It is a tree of statements, not of blocks. Loop is a node with a body, the way it is in the source. A conventional optimising compiler would have destroyed this by now — lowered the loop to a comparison, a branch and a back-edge over basic blocks in SSA form, which is the right representation for register allocation and the wrong one for printing do k = 1, nlayers back out. PSyIR’s job is to survive a round trip through a human, so it keeps the structure the human wrote.3
A Loop has exactly four children: start, stop, step, body — in that order, positionally. Not a bounds object, not keywords.
Schedule is the body. It exists so that a loop body is a node, and therefore something a transformation can wrap a directive around. When you see Schedule[] in a tree view, read “block”.
Nothing is typed yet, quite. Literal[value:'1', Scalar<INTEGER, UNDEFINED>] — the kind is UNDEFINED because the literal 1 has no kind suffix, while 2.0_r_def carries Scalar<REAL, Reference[name:'r_def']>: the kind is a reference to a symbol, not a number, because PSyclone has not read constants_mod and does not know what r_def is. It does not need to.
Every scope carries a symbol table, and it is the part of PSyIR you will actually query in a script — “is this name a local, an argument, or an import?” decides whether a variable can be made private.
from psyclone.psyir.nodes import Routine
routine = psyir(source).walk(Routine)[0]
table = routine.symbol_table
for symbol in table.symbols:
print(f"{symbol.name:8s} {type(symbol.interface).__name__:20s} {symbol.datatype}")
print("\nargument order:", [s.name for s in table.argument_list])
print("r_def is:", table.lookup("r_def").interface)n ArgumentInterface Scalar<INTEGER, Reference[name:'i_def']>
nlayers ArgumentInterface Scalar<INTEGER, Reference[name:'i_def']>
kappa ArgumentInterface Scalar<REAL, Reference[name:'r_def']>
field ArgumentInterface Array<Scalar<REAL, Reference[name:'r_def']>, shape=[Reference[name:'n'], Reference[name:'nlayers']]>
i AutomaticInterface Scalar<INTEGER, Reference[name:'i_def']>
k AutomaticInterface Scalar<INTEGER, Reference[name:'i_def']>
argument order: ['field', 'kappa', 'n', 'nlayers']
r_def is: Import(container='constants_mod')Three interfaces, three meanings. ArgumentInterface — a dummy argument, and the table’s argument_list preserves the order, which is how PSyclone knows how to call the routine. AutomaticInterface — an ordinary local, and therefore the class of thing that can be privatised for OpenMP. ImportInterface — brought in by use, so it lives in another file that PSyclone may never have read; r_def resolves to Import(container='constants_mod') and nothing more.
That last one is why LFRic’s kind system is invisible to PSyclone, and why that is fine. real(r_def) is Scalar<REAL, Reference['r_def']>: enough to know it is a real, enough to print it back out identically, not enough to know it is 64-bit. Almost no transformation needs to know.
CodeBlock: the most important nodeHere is the one that will actually bite you. PSyIR does not model all of Fortran. When the frontend meets something it has no node for, it does not fail — it wraps the raw fparser2 subtree in a CodeBlock and carries on.
from psyclone.psyir.nodes import CodeBlock
mixed = psyir("""
subroutine report(a, n)
integer, intent(in) :: n
real, intent(inout) :: a(n)
integer :: i
do i = 1, n
a(i) = a(i) * 2.0
end do
write(*,*) 'total ', sum(a)
end subroutine report
""")
for block in mixed.walk(CodeBlock):
print(block.structure, "->", block.get_fortran_lines())Structure.STATEMENT -> ["WRITE(*, *) 'total ', SUM(a)"]The loop is fully modelled; the write is a black box that PSyclone can move but not look inside. This is a good trade — it is what lets a 943-line physics routine be transformed when one statement in it is exotic — but it has a hard consequence: a CodeBlock inside a region blocks transformations of that region.
from psyclone.psyir.nodes import Loop
from psyclone.psyir.transformations import OMPParallelLoopTrans, TransformationError
blocked = psyir("""
subroutine noisy(a, n)
integer, intent(in) :: n
real, intent(inout) :: a(n)
integer :: i
do i = 1, n
a(i) = a(i) * 2.0
write(*,*) i
end do
end subroutine noisy
""")
try:
OMPParallelLoopTrans().apply(blocked.walk(Loop)[0])
except TransformationError as err:
print(err)Transformation Error: Nodes of type 'Fparser2CodeBlock' cannot be enclosed by a OMPParallelLoopTrans transformation (use the 'node-type-check: False' option to accept them at your own risk) in:
do i = 1, n, 1
a(i) = a(i) * 2.0
! PSyclone CodeBlock (unsupported code) reason:
! - Unsupported statement: Write_Stmt
WRITE(*, *) i
enddo
PSyclone cannot know what that write touches, so it will not certify the loop as parallel. The message names the escape hatch (node-type-check: False) and you should almost never take it: here you would be promising that a write from many threads at once is fine, which it is not.
This is the first thing to check when a transformation mysteriously does nothing in a real routine. Print the CodeBlocks, find out which statement lost you the loop, and consider whether the source should change — which is frequently the right answer, and is part of what ticket #813 did.
Round-tripping is lossy, on purpose, and you need to know exactly how, because you will diff generated code against generated code for the rest of your career.
module diffuse_mod
use constants_mod, only : i_def, r_def
implicit none
public
contains
subroutine diffuse(field, kappa, n, nlayers)
integer(kind=i_def), intent(in) :: n
integer(kind=i_def), intent(in) :: nlayers
real(kind=r_def), intent(in) :: kappa
real(kind=r_def), dimension(n,nlayers), intent(inout) :: field
integer(kind=i_def) :: i
integer(kind=i_def) :: k
do k = 1, nlayers, 1
do i = 2, n - 1, 1
field(i,k) = field(i,k) + kappa * (field(i - 1,k) - 2.0_r_def * field(i,k) + field(i + 1,k))
enddo
enddo
end subroutine diffuse
end module diffuse_modCompare with the input. Declarations are split one per line; the only : list is alphabetised; dimension(n,nlayers) replaces field(n,nlayers); intent(in) is spelled out on each; do k = 1, nlayers gains an explicit , 1 step; end do becomes enddo; public is stated; blank lines are normalised. All semantically identical, all diff noise.
And the big one:
gw_source = LFRIC_APPS / "science/physics_schemes/source/gravity_wave_drag/gw_ussp_mod.F90"
run("psyclone", "-l", "all", "-o", "/tmp/gw_plain.f90", gw_source)
run("psyclone", "-l", "all", "--keep-comments", "--keep-directives",
"-o", "/tmp/gw_commented.f90", gw_source)
def comment_count(path):
return sum(1 for line in open(path) if line.lstrip().startswith("!"))
print("source comments :", comment_count(gw_source))
print("default output :", comment_count("/tmp/gw_plain.f90"))
print("--keep-comments :", comment_count("/tmp/gw_commented.f90"))
source comments : 377
default output : 0
--keep-comments : 365By default every comment is deleted. 377 of them, from a routine whose comments are the only documentation of a fairly deep piece of gravity-wave physics. --keep-comments and --keep-directives were added in 3.3 and recover almost all of them; LFRic’s psyclone_transmute.mk does not yet pass either flag, which is a small and worthwhile patch.4
Two more flags worth knowing now. -l all limits output lines to 132 characters — LFRic passes it everywhere, because PSyclone’s writer happily emits a 400-character line that no Fortran compiler will accept in free form. -o names the output file; without it, output goes to stdout, which is what makes PSyclone pleasant to explore from a shell.
A transformation is an object with two methods: validate, which raises TransformationError if the rewrite would be illegal, and apply, which validates and then mutates the tree in place. That is the whole interface, and the asymmetry in how much you care about each is worth stating: apply is the boring half.
one_loop = psyir("""
subroutine scale_column(field, kappa, n, nlayers)
integer, intent(in) :: n, nlayers
real, intent(in) :: kappa
real, intent(inout) :: field(n, nlayers)
integer :: i, k
do k = 1, nlayers
do i = 1, n
field(i,k) = kappa * field(i,k)
end do
end do
end subroutine scale_column
""")
before = fortran(one_loop)
OMPParallelLoopTrans(omp_schedule="static").apply(one_loop.walk(Loop)[0])
diff(before, fortran(one_loop))Note what PSyclone worked out for itself: private(i,k) — including k, the loop variable of the parallelised loop, and i, the inner one — and default(shared). Getting a private list wrong by hand is the classic OpenMP bug, silent and load-dependent, and deriving it from the tree is most of the value on offer here.
Transformation Error: Loop cannot be parallelised because:
Error: The write access to 'a(i)' in 'a(i) = a(i - 1) + 1.0' and the read access to 'a(i - 1)' in 'a(i) = a(i - 1) + 1.0' are dependent and cannot be parallelised. Variable: 'a'.
Consider using the "ignore_dependencies_for" transformation option if this is a false dependency
Consider using the "array_privatisation" transformation option if this is a write-write dependencyA loop-carried dependence, found by comparing the subscripts of the write and the read, and reported with both. sed would have parallelised this and produced a program that gives a different answer on every run.
The two hints in that message are the ones you will reach for on real code, and they are promises, not requests:
ignore_dependencies_for=[...] — “this dependence is spurious, take my word for it”. Usually because the real index relationship is beyond PSyclone’s analysis (an indirection through a dofmap, say).array_privatisation — for a write-write dependence on scratch space, where each iteration wants its own copy.Getting either wrong buys you a race condition, checked into a model that runs on 384 ranks. The dependence analysis is also available directly, which is how a script can make its own decisions:
a: READ+WRITE, i: READ. That map — signature to access type — is the primitive under most of the LFRic helper functions in §6.
two_passes = """
subroutine two_passes(a, b, n)
integer, intent(in) :: n
real, intent(inout) :: a(n), b(n)
integer :: i
do i = 1, n
a(i) = a(i) * 2.0
end do
do i = 1, n
b(i) = b(i) + a(i)
end do
end subroutine two_passes
"""
from psyclone.psyir.transformations import LoopFuseTrans
fused = psyir(two_passes)
before = fortran(fused)
loops = fused.walk(Loop)
LoopFuseTrans().apply(loops[0], loops[1])
diff(before, fortran(fused))One pass over memory instead of two — the same win as fusing two NumPy expressions, for the same reason, and worth the same amount: on a memory-bound stencil, most of it.
Now shift one index and ask again:
Transformation Error: LoopFuseTrans. Error: Variable 'a' is used with different indices: 'a(i)' and 'a(i + 1)'. Variable: 'a'.Correct: after fusion, iteration i would read a(i+1) before the first loop’s iteration i+1 had written it. The analysis is deliberately conservative — “used with different indices” is a syntactic test, and it will refuse fusions that happen to be legal. That is the right bias for a tool whose mistakes get committed.
The one genuinely annoying thing about the API. There are three transformation modules and the split is historical, not logical:
| module | holds | notes |
|---|---|---|
psyclone.psyir.transformations |
the generic ones: OMPParallelTrans, LoopFuseTrans, LoopSwapTrans, InlineTrans, ChunkLoopTrans, ProfileTrans |
where new work lands |
psyclone.transformations |
the top-level re-export, plus the PSyKAl-aware ones: LFRicColourTrans, LFRicOMPLoopTrans, ACCParallelTrans |
re-exports shift between releases |
psyclone.domain.lfric.transformations |
LFRic-specific: LFRicRedundantComputationTrans, LFRicLoopFuseTrans, LFRicExtractTrans |
Some names are importable from two of them, which is exactly how a script ends up working on one release and failing on the next. In 3.3, OMPParallelTrans left psyclone.transformations, and Dynamo0p3ColourTrans/Dynamo0p3OMPLoopTrans/Dynamo0p3RedundantComputationTrans became LFRicColourTrans/LFRicOMPLoopTrans/LFRicRedundantComputationTrans.5 When an LFRic optimisation script raises ImportError, that is the first thing to check, and the fix is mechanical.
Now the DSL, with the IR already in hand.
PSyKAl = Parallel System, Kernel, Algorithm: three layers with a rule about what each may contain. The rule is the entire idea.
| layer | file | written by | may contain | may not contain |
|---|---|---|---|---|
| Algorithm | .x90 |
a scientist | whole fields, call invoke(...) |
loops over cells, indices, MPI, OpenMP |
| PSy | generated _psy.f90 |
PSyclone | loops, dofmaps, halo exchanges, directives | science |
| Kernel | _kernel_mod.F90 |
a scientist | one column’s arithmetic, intrinsic types | fields, use of anything global, I/O |
The middle layer is the only one anybody generates, and it is where every decision about how the model runs is expressed. The two human layers are written once and retargeted for free — that is the bet.
type, public, extends(kernel_type) :: matrix_vector_kernel_type
private
type(arg_type) :: meta_args(3) = (/ &
arg_type(GH_FIELD, GH_REAL, GH_INC, ANY_SPACE_1), &
arg_type(GH_FIELD, GH_REAL, GH_READ, ANY_SPACE_2), &
arg_type(GH_OPERATOR, GH_REAL, GH_READ, ANY_SPACE_1, ANY_SPACE_2) &
/)
integer :: operates_on = CELL_COLUMN
end type
!-------------------------------------------------------------------------------
! Contained functions/subroutines
!-------------------------------------------------------------------------------
public :: matrix_vector_code
! Generic interface for real32 and real64 types
interface matrix_vector_code
module procedure &
matrix_vector_code_r_single, &
matrix_vector_code_r_double
end interfaceThe declaration is the interface. meta_args is an array of arg_type in the order the algorithm layer passes them, each naming three or four things: what kind of argument (GH_FIELD, GH_OPERATOR, GH_SCALAR), of what data type (GH_REAL), accessed how (GH_READ, GH_INC, GH_WRITE, GH_READWRITE), on which function space (W1, W3, ANY_SPACE_1). Then operates_on says what one call is handed: CELL_COLUMN for almost everything, DOF for pointwise work, DOMAIN for the rare kernel that wants the lot.
The access descriptor is the load-bearing part, and it is the same information as a JAX donation annotation, doing more work:
| descriptor | means | what PSyclone must then do |
|---|---|---|
GH_READ |
read only | may need a halo exchange first |
GH_WRITE |
overwrite, no read | nothing; safe to thread |
GH_READWRITE |
read and write, own dofs only | discontinuous spaces only |
GH_INC |
increment a shared dof | colour the mesh, or race; and loop into the halo |
GH_INC on a continuous function space is the interesting case, and the source of most of the generated complexity below: continuous means neighbouring cells share degrees of freedom, so two threads on adjacent cells increment the same array entry.
demo/ holds a self-contained example built for exactly this — a kernel with GH_INC on W1, an algorithm with two invokes, and an optimisation script. Start with no script at all:
MODULE mini_alg_mod
USE constants_mod, ONLY: r_def, i_def
USE field_mod, ONLY: field_type
USE function_space_mod, ONLY: function_space_type
IMPLICIT NONE
PRIVATE
PUBLIC :: mini_alg
CONTAINS
SUBROUTINE mini_alg(lhs, rhs)
USE mini_alg_mod_psy, ONLY: invoke_1
USE mini_alg_mod_psy, ONLY: invoke_apply_mini_matvec
IMPLICIT NONE
TYPE(field_type), INTENT(INOUT) :: lhs
TYPE(field_type), INTENT(IN) :: rhs
CALL invoke_apply_mini_matvec(lhs, rhs)
CALL invoke_1(lhs)
END SUBROUTINE mini_alg
END MODULE mini_alg_modThe algorithm layer has been rewritten: each invoke became a call to a generated subroutine, and a use of a module that did not exist a second ago. Note invoke_apply_mini_matvec took its name from the name= argument in the source, while the anonymous second invoke became invoke_1 — a good reason to always name invokes, since these are the symbols that will appear in a profile.
And the layer that call lands in:
module mini_alg_mod_psy
use field_mod, only : field_proxy_type, field_proxy_type_1=>field_proxy_type, field_type
use mini_matvec_kernel_mod, only : mini_matvec_code
use constants_mod, only : r_def
implicit none
public
contains
subroutine invoke_apply_mini_matvec(lhs, rhs)
use mesh_mod, only : mesh_type
use constants_mod, only : i_def
type(field_type), intent(in) :: lhs
type(field_type), intent(in) :: rhs
integer(kind=i_def) :: df
integer(kind=i_def) :: cell
type(mesh_type), pointer :: mesh => null()
integer(kind=i_def) :: max_halo_depth_mesh
real(kind=r_def), pointer, dimension(:) :: lhs_data => null()
real(kind=r_def), pointer, dimension(:) :: rhs_data => null()
integer(kind=i_def) :: nlayers_lhs
integer(kind=i_def) :: ndf_as1_lhs
integer(kind=i_def) :: undf_as1_lhs
integer(kind=i_def) :: ndf_w1
integer(kind=i_def) :: undf_w1
integer(kind=i_def) :: ndf_w3
integer(kind=i_def) :: undf_w3
integer(kind=i_def), pointer :: map_w1(:,:) => null()
integer(kind=i_def), pointer :: map_w3(:,:) => null()
type(field_proxy_type) :: lhs_proxy
type(field_proxy_type) :: rhs_proxy
integer(kind=i_def) :: loop0_start
integer(kind=i_def) :: loop0_stop
integer(kind=i_def) :: loop1_start
integer(kind=i_def) :: loop1_stop
! Initialise field and/or operator proxies
lhs_proxy = lhs%get_proxy()
lhs_data => lhs_proxy%data
rhs_proxy = rhs%get_proxy()
rhs_data => rhs_proxy%data
! Initialise number of layers
nlayers_lhs = lhs_proxy%vspace%get_nlayers()
! Create a mesh object
mesh => lhs_proxy%vspace%get_mesh()
max_halo_depth_mesh = mesh%get_halo_depth()
! Look-up dofmaps for each function space
map_w1 => lhs_proxy%vspace%get_whole_dofmap()
map_w3 => rhs_proxy%vspace%get_whole_dofmap()
! Initialise number of DoFs for as1_lhs
ndf_as1_lhs = lhs_proxy%vspace%get_ndf()
undf_as1_lhs = lhs_proxy%vspace%get_undf()
! Initialise number of DoFs for w1
ndf_w1 = lhs_proxy%vspace%get_ndf()
undf_w1 = lhs_proxy%vspace%get_undf()
! Initialise number of DoFs for w3
ndf_w3 = rhs_proxy%vspace%get_ndf()
undf_w3 = rhs_proxy%vspace%get_undf()
! Set-up all of the loop bounds
loop0_start = 1
loop0_stop = lhs_proxy%vspace%get_last_dof_annexed()
loop1_start = 1
loop1_stop = mesh%get_last_halo_cell(1)
! Call kernels and communication routines
do df = loop0_start, loop0_stop, 1
! Built-in: setval_c (set a real-valued field to a real scalar value)
lhs_data(df) = 0.0_r_def
enddo
! Set halos dirty/clean for fields modified in the above loop(s)
call lhs_proxy%set_dirty()
if (rhs_proxy%is_dirty(depth=1)) then
call rhs_proxy%halo_exchange(depth=1)
end if
do cell = loop1_start, loop1_stop, 1
call mini_matvec_code(nlayers_lhs, lhs_data, rhs_data, ndf_w1, undf_w1, map_w1(:,cell), ndf_w3, undf_w3, map_w3(:,cell))
enddo
! Set halos dirty/clean for fields modified in the above loop(s)
call lhs_proxy%set_dirty()
end subroutine invoke_apply_mini_matvec
subroutine invoke_1(lhs)
use mesh_mod, only : mesh_type
use constants_mod, only : i_def
type(field_type), intent(in) :: lhs
integer(kind=i_def) :: df
type(mesh_type), pointer :: mesh => null()
integer(kind=i_def) :: max_halo_depth_mesh
real(kind=r_def), pointer, dimension(:) :: lhs_data => null()
type(field_proxy_type_1) :: lhs_proxy
integer(kind=i_def) :: loop0_start
integer(kind=i_def) :: loop0_stop
! Initialise field and/or operator proxies
lhs_proxy = lhs%get_proxy()
lhs_data => lhs_proxy%data
! Create a mesh object
mesh => lhs_proxy%vspace%get_mesh()
max_halo_depth_mesh = mesh%get_halo_depth()
! Set-up all of the loop bounds
loop0_start = 1
loop0_stop = lhs_proxy%vspace%get_last_dof_annexed()
! Call kernels and communication routines
do df = loop0_start, loop0_stop, 1
! Built-in: inc_a_times_X (scale a real-valued field)
lhs_data(df) = 2.0_r_def * lhs_data(df)
enddo
! Set halos dirty/clean for fields modified in the above loop(s)
call lhs_proxy%set_dirty()
end subroutine invoke_1
end module mini_alg_mod_psyRead it in order, because it is the same order every generated PSy routine follows:
lhs_proxy = lhs%get_proxy() then lhs_data => lhs_proxy%data. The proxy pattern from the Fortran note: field_type keeps its data private, the PSy layer needs a raw array, and get_proxy is the one sanctioned hole through the encapsulation.nlayers, ndf_w1 (dofs per cell), undf_w1 (unique dofs on this partition), map_w1 => ...get_whole_dofmap().loop1_stop = mesh%get_last_halo_cell(1). Not ncells — the loop runs into the level-1 halo, because GH_INC means owned dofs on the partition boundary are only complete once the neighbouring halo cell has contributed too.if (rhs_proxy%is_dirty(depth=1)) then guarding call rhs_proxy%halo_exchange(depth=1) — inserted because the kernel reads rhs in cells that include halo cells, and guarded so that a clean halo costs nothing. Nobody wrote that; it was derived from GH_READ plus the loop bound.map_w1(:,cell) — the dofmap column for this cell.call lhs_proxy%set_dirty(), because the loop wrote lhs and its halo copies elsewhere are now stale.That is the whole distributed-memory protocol, generated from three lines of metadata. The two setval_c/inc_a_times_X loops have no kernel file anywhere — they are builtins, whose bodies PSyclone writes itself.
Now with demo/optimisation/global.py, which colours continuous-space cell loops and then threads them:
# -----------------------------------------------------------------------------
# A teaching-sized version of lfric_core's
# infrastructure/build/psyclone/psyclone_tools.py, laid out so that the three
# decisions are visible in twenty lines rather than two hundred.
#
# The file name and location are not free. The LFRic build system looks for
# `<optimisation-path>/psykal/global.py` (a whole-application script) or
# `<optimisation-path>/psykal/<algorithm-name>.py` (a per-file override), and
# passes whichever it finds to `psyclone -s`. PSyclone then imports the file
# and calls the single function `trans` with the PSyIR of the PSy layer.
# -----------------------------------------------------------------------------
"""Colour, then parallelise, the loops in the generated PSy layer."""
from psyclone.domain.lfric import LFRicConstants
from psyclone.psyGen import InvokeSchedule
from psyclone.psyir.nodes import Directive, Loop
from psyclone.psyir.transformations import OMPParallelTrans
from psyclone.transformations import LFRicColourTrans, LFRicOMPLoopTrans
def trans(psyir):
"""Entry point. PSyclone calls this with the PSy layer's PSyIR.
:param psyir: the PSyIR of the PSy-layer.
:type psyir: :py:class:`psyclone.psyir.nodes.FileContainer`
"""
const = LFRicConstants()
colour = LFRicColourTrans()
omp_region = OMPParallelTrans()
omp_loop = LFRicOMPLoopTrans()
for schedule in psyir.walk(InvokeSchedule):
# 1. Colour every loop over cells on a *continuous* function space.
# Discontinuous spaces (W3, Wtheta) need no colouring -- no two
# cells share a dof, so there is no race to avoid -- and loops
# over dofs are not over cells at all.
for loop in schedule.walk(Loop):
if (
loop.iteration_space.endswith("cell_column")
and loop.field_space.orig_name
not in const.VALID_DISCONTINUOUS_NAMES
):
colour.apply(loop)
# 2. Parallelise. After colouring, the loop nest is
# `do colour / do cell-in-colour`, and it is the *inner* loop that
# is safe to thread -- which is why the `loop_type == "colours"`
# outer loop is skipped rather than parallelised.
for loop in schedule.loops():
if loop.loop_type not in ("colours", "null") and not loop.ancestor(
Directive
):
omp_region.apply(loop)
omp_loop.apply(loop, options={"reprod": True})run("psyclone", "-api", "lfric", "--config", LFRIC_CONFIG,
"-d", "demo/kernel", "-s", "./demo/optimisation/global.py",
"-oalg", "/tmp/mini_alg_opt.f90", "-opsy", "/tmp/mini_psy_opt.f90",
"demo/algorithm/mini_alg_mod.x90")
diff(open("/tmp/mini_psy.f90").read(), open("/tmp/mini_psy_opt.f90").read(),
"psy layer, no script", "psy layer, colour + omp")--- psy layer, no script
+++ psy layer, colour + omp
@@ -18,6 +18,10 @@
real(kind=r_def), pointer, dimension(:) :: lhs_data => null()
real(kind=r_def), pointer, dimension(:) :: rhs_data => null()
integer(kind=i_def) :: nlayers_lhs
+ integer(kind=i_def) :: colour
+ integer(kind=i_def), pointer :: cmap(:,:)
+ integer(kind=i_def) :: ncolour
+ integer(kind=i_def), allocatable, dimension(:,:) :: last_halo_cell_all_colours
integer(kind=i_def) :: ndf_as1_lhs
integer(kind=i_def) :: undf_as1_lhs
integer(kind=i_def) :: ndf_w1
@@ -32,6 +36,7 @@
integer(kind=i_def) :: loop0_stop
integer(kind=i_def) :: loop1_start
integer(kind=i_def) :: loop1_stop
+ integer(kind=i_def) :: loop2_start
! Initialise field and/or operator proxies
lhs_proxy = lhs%get_proxy()
@@ -45,6 +50,10 @@
! Create a mesh object
mesh => lhs_proxy%vspace%get_mesh()
max_halo_depth_mesh = mesh%get_halo_depth()
+
+ ! Get the colourmap
+ ncolour = mesh%get_ncolours()
+ cmap => mesh%get_colour_map()
! Look-up dofmaps for each function space
map_w1 => lhs_proxy%vspace%get_whole_dofmap()
@@ -61,26 +70,38 @@
! Initialise number of DoFs for w3
ndf_w3 = rhs_proxy%vspace%get_ndf()
undf_w3 = rhs_proxy%vspace%get_undf()
+ last_halo_cell_all_colours = mesh%get_last_halo_cell_all_colours()
! Set-up all of the loop bounds
loop0_start = 1
loop0_stop = lhs_proxy%vspace%get_last_dof_annexed()
loop1_start = 1
- loop1_stop = mesh%get_last_halo_cell(1)
+ loop1_stop = ncolour
+ loop2_start = 1
! Call kernels and communication routines
+ !$omp parallel default(shared) private(df)
+ !$omp do schedule(static)
do df = loop0_start, loop0_stop, 1
! Built-in: setval_c (set a real-valued field to a real scalar value)
lhs_data(df) = 0.0_r_def
enddo
+ !$omp end do
+ !$omp end parallel
! Set halos dirty/clean for fields modified in the above loop(s)
call lhs_proxy%set_dirty()
if (rhs_proxy%is_dirty(depth=1)) then
call rhs_proxy%halo_exchange(depth=1)
end if
- do cell = loop1_start, loop1_stop, 1
- call mini_matvec_code(nlayers_lhs, lhs_data, rhs_data, ndf_w1, undf_w1, map_w1(:,cell), ndf_w3, undf_w3, map_w3(:,cell))
+ do colour = loop1_start, loop1_stop, 1
+ !$omp parallel default(shared) private(cell)
+ !$omp do schedule(static)
+ do cell = loop2_start, last_halo_cell_all_colours(colour,1), 1
+ call mini_matvec_code(nlayers_lhs, lhs_data, rhs_data, ndf_w1, undf_w1, map_w1(:,cmap(colour,cell)), ndf_w3, undf_w3, map_w3(:,cmap(colour,cell)))
+ enddo
+ !$omp end do
+ !$omp end parallel
enddo
! Set halos dirty/clean for fields modified in the above loop(s)
@@ -112,10 +133,14 @@
loop0_stop = lhs_proxy%vspace%get_last_dof_annexed()
! Call kernels and communication routines
+ !$omp parallel default(shared) private(df)
+ !$omp do schedule(static)
do df = loop0_start, loop0_stop, 1
! Built-in: inc_a_times_X (scale a real-valued field)
lhs_data(df) = 2.0_r_def * lhs_data(df)
enddo
+ !$omp end do
+ !$omp end parallel
! Set halos dirty/clean for fields modified in the above loop(s)
call lhs_proxy%set_dirty()The shape to recognise:
Colours outside, threads inside, and cell indices reached through cmap. The mesh has been partitioned so that no two cells of the same colour share a degree of freedom, so within one colour the increments cannot collide; the implicit barrier at the end of each !$omp parallel do separates colour from colour. This is chromatic Gibbs sampling with the conditional-independence graph replaced by the dof-sharing graph — the same algorithm, the same reason, and if you have written a blocked Gibbs sampler you have already had this idea.6
Note the price, which is why colouring is applied only where required: cmap(colour,cell) is an indirection, so cells visited consecutively are no longer contiguous in memory, and the prefetcher suffers. On a discontinuous space (W3, Wtheta) no dof is shared, no colouring is needed, and the script must not add it — which is exactly the VALID_DISCONTINUOUS_NAMES test in the script.
Everything above was a miniature. Here is the actual production script — lfric_core/infrastructure/build/psyclone/psyclone_tools.py, imported by every application’s global.py — run over a real algorithm from lfric_core:
alg = LFRIC_CORE / "applications/simple_diffusion/source/algorithm/simple_diffusion_alg_mod.x90"
script = LFRIC_CORE / "applications/simple_diffusion/optimisation/meto-ex1a/psykal/global.py"
run("psyclone", "-api", "lfric", "--config", LFRIC_CONFIG,
"-d", LFRIC_CORE / "applications/simple_diffusion/source",
"-s", script,
"-oalg", "/tmp/sd_alg.f90", "-opsy", "/tmp/sd_psy.f90", alg,
env={"PYTHONPATH": str(LFRIC_PSYCLONE_TOOLS)})Transformed invoke 'invoke_compute_diffusion' ... Found 2 ['setval_c'] calls InvokeSchedule[invoke='invoke_compute_diffusion', dm=True] 0: OMPParallelDirective[] Schedule[] 0: OMPDoDirective[omp_schedule=static] Schedule[] 0: Loop[type='dof', field_space='any_space_1', it_space='dof', upper_bound='dof_halo(Literal[value:'1', Scalar<INTEGER, UNDEFINED>])'] Reference[name:'uninitialised_loop0_start'] Reference[name:'uninitialised_loop0_stop'] Literal[value:'1', Scalar<INTEGER, UNDEFINED>] Schedule[] 0: BuiltIn setval_c(visc,visc_val) OMPDefaultClause[default=DefaultClauseTypes.SHARED] OMPPrivateClause[] OMPFirstprivateClause[] 1: OMPParallelDirective[] Schedule[] 0: OMPDoDirective[omp_schedule=static] Schedule[] 0: Loop[type='dof', field_space='any_space_1', it_space='dof', upper_bound='dof_halo(Literal[value:'1', Scalar<INTEGER, UNDEFINED>])'] Reference[name:'uninitialised_loop1_start'] Reference[name:'uninitialised_loop1_stop'] Literal[value:'1', Scalar<INTEGER, UNDEFINED>] Schedule[] 0: BuiltIn setval_c(dfield_in,0.0_r_def) OMPDefaultClause[default=DefaultClauseTypes.SHARED] OMPPrivateClause[] OMPFirstprivateClause[] 2: HaloExchange[field='field_in', type='cross', depth=stencil_depth, check_dirty=True] 3: OMPParallelDirective[] Schedule[] 0: OMPDoDirective[omp_schedule=static] Schedule[] 0: Loop[type='', field_space='wtheta', it_space='cell_column', upper_bound='ncells'] Reference[name:'uninitialised_loop2_start'] Reference[name:'uninitialised_loop2_stop'] Literal[value:'1', Scalar<INTEGER, UNDEFINED>] Schedule[] 0: CodedKern tracer_tutorial_diff_code(dfield_in,field_in,visc,dx_at_w2) [module_inline=False] OMPDefaultClause[default=DefaultClauseTypes.SHARED] OMPPrivateClause[] OMPFirstprivateClause[] Transformed invoke 'invoke_1' ... Found 2 ['setval_c'] calls InvokeSchedule[invoke='invoke_1', dm=True] 0: OMPParallelDirective[] Schedule[] 0: OMPDoDirective[omp_schedule=static] Schedule[] 0: Loop[type='dof', field_space='any_space_1', it_space='dof', upper_bound='nannexed'] Reference[name:'uninitialised_loop0_start'] Reference[name:'uninitialised_loop0_stop'] Literal[value:'1', Scalar<INTEGER, UNDEFINED>] Schedule[] 0: BuiltIn inc_x_plus_y(field_in,dfield_in) OMPDefaultClause[default=DefaultClauseTypes.SHARED] OMPPrivateClause[] OMPFirstprivateClause[]
"Transformed invoke 'invoke_compute_diffusion' ...\nFound 2 ['setval_c'] calls\n\x1b[97mInvokeSchedule\x1b[0m[invoke='invoke_compute_diffusion', dm=True]\n 0: \x1b[32mOMPParallelDirective\x1b[0m[]\n \x1b[97mSchedule\x1b[0m[]\n 0: \x1b[32mOMPDoDirective\x1b[0m[omp_schedule=static]\n \x1b[97mSchedule\x1b[0m[]\n 0: \x1b[31mLoop\x1b[0m[type='dof', field_space='any_space_1', it_space='dof', upper_bound='dof_halo(Literal[value:'1', Scalar<INTEGER, UNDEFINED>])']\n \x1b[33mReference\x1b[0m[name:'uninitialised_loop0_start']\n \x1b[33mReference\x1b[0m[name:'uninitialised_loop0_stop']\n \x1b[33mLiteral\x1b[0m[value:'1', Scalar<INTEGER, UNDEFINED>]\n \x1b[97mSchedule\x1b[0m[]\n 0: \x1b[35mBuiltIn\x1b[0m setval_c(visc,visc_val)\n \x1b[32mOMPDefaultClause\x1b[0m[default=DefaultClauseTypes.SHARED]\n \x1b[32mOMPPrivateClause\x1b[0m[]\n \x1b[32mOMPFirstprivateClause\x1b[0m[]\n 1: \x1b[32mOMPParallelDirective\x1b[0m[]\n \x1b[97mSchedule\x1b[0m[]\n 0: \x1b[32mOMPDoDirective\x1b[0m[omp_schedule=static]\n \x1b[97mSchedule\x1b[0m[]\n 0: \x1b[31mLoop\x1b[0m[type='dof', field_space='any_space_1', it_space='dof', upper_bound='dof_halo(Literal[value:'1', Scalar<INTEGER, UNDEFINED>])']\n \x1b[33mReference\x1b[0m[name:'uninitialised_loop1_start']\n \x1b[33mReference\x1b[0m[name:'uninitialised_loop1_stop']\n \x1b[33mLiteral\x1b[0m[value:'1', Scalar<INTEGER, UNDEFINED>]\n \x1b[97mSchedule\x1b[0m[]\n 0: \x1b[35mBuiltIn\x1b[0m setval_c(dfield_in,0.0_r_def)\n \x1b[32mOMPDefaultClause\x1b[0m[default=DefaultClauseTypes.SHARED]\n \x1b[32mOMPPrivateClause\x1b[0m[]\n \x1b[32mOMPFirstprivateClause\x1b[0m[]\n 2: \x1b[34mHaloExchange\x1b[0m[field='field_in', type='cross', depth=stencil_depth, check_dirty=True]\n 3: \x1b[32mOMPParallelDirective\x1b[0m[]\n \x1b[97mSchedule\x1b[0m[]\n 0: \x1b[32mOMPDoDirective\x1b[0m[omp_schedule=static]\n \x1b[97mSchedule\x1b[0m[]\n 0: \x1b[31mLoop\x1b[0m[type='', field_space='wtheta', it_space='cell_column', upper_bound='ncells']\n \x1b[33mReference\x1b[0m[name:'uninitialised_loop2_start']\n \x1b[33mReference\x1b[0m[name:'uninitialised_loop2_stop']\n \x1b[33mLiteral\x1b[0m[value:'1', Scalar<INTEGER, UNDEFINED>]\n \x1b[97mSchedule\x1b[0m[]\n 0: \x1b[35mCodedKern\x1b[0m tracer_tutorial_diff_code(dfield_in,field_in,visc,dx_at_w2) [module_inline=False]\n \x1b[32mOMPDefaultClause\x1b[0m[default=DefaultClauseTypes.SHARED]\n \x1b[32mOMPPrivateClause\x1b[0m[]\n \x1b[32mOMPFirstprivateClause\x1b[0m[]\n\nTransformed invoke 'invoke_1' ...\nFound 2 ['setval_c'] calls\n\x1b[97mInvokeSchedule\x1b[0m[invoke='invoke_1', dm=True]\n 0: \x1b[32mOMPParallelDirective\x1b[0m[]\n \x1b[97mSchedule\x1b[0m[]\n 0: \x1b[32mOMPDoDirective\x1b[0m[omp_schedule=static]\n \x1b[97mSchedule\x1b[0m[]\n 0: \x1b[31mLoop\x1b[0m[type='dof', field_space='any_space_1', it_space='dof', upper_bound='nannexed']\n \x1b[33mReference\x1b[0m[name:'uninitialised_loop0_start']\n \x1b[33mReference\x1b[0m[name:'uninitialised_loop0_stop']\n \x1b[33mLiteral\x1b[0m[value:'1', Scalar<INTEGER, UNDEFINED>]\n \x1b[97mSchedule\x1b[0m[]\n 0: \x1b[35mBuiltIn\x1b[0m inc_x_plus_y(field_in,dfield_in)\n \x1b[32mOMPDefaultClause\x1b[0m[default=DefaultClauseTypes.SHARED]\n \x1b[32mOMPPrivateClause\x1b[0m[]\n \x1b[32mOMPFirstprivateClause\x1b[0m[]\n\n"That output is the script’s own view_transformed_schedule talking — production LFRic scripts print the transformed schedule into the build log, which is how you find out what happened to your loops without reading 400 lines of generated Fortran. Worth reading closely:
HaloExchange[field='field_in', type='cross', depth=stencil_depth, check_dirty=True] — a stencil exchange, depth taken from a variable in the algorithm layer, guarded by a dirty check.Loop[type='dof', it_space='dof', upper_bound='nannexed'] for the builtins, versus it_space='cell_column' for the kernel.upper_bound='nannexed': annexed dofs are dofs owned by another rank but sitting on cells this rank touches. Whether they are computed is a config file setting (COMPUTE_ANNEXED_DOFS), not a script setting — which is one reason --config is not optional for the LFRic API.The three passes in psyclone_tools.py are worth naming since you will read them in every ticket:
redundant_computation_setval extends setval_* builtins to compute into the level-1 halo. Deliberately doing extra arithmetic to avoid a communication — because the alternative is a halo exchange, and on 384 ranks arithmetic is free and messages are not. colour_loops as above. openmp_parallelise_loops wraps each remaining loop in a parallel region and adds !$omp do, with reprod=True — reproducible reductions, meaning a fixed summation order at some cost in speed, so that the model gives bit-identical answers on the same rank count. In a model validated by checksums, that is not negotiable.
The ordering constraints between them are real, and psyclone_tools.py enforces them by raising: colour before profile, profile before OpenMP. A ProfileNode inserted first would be in the way of the colouring transformation, and a profiling caliper inside a parallel region measures something other than what you meant.
psyclone-kern: checking a kernel against its own metadataA kernel’s argument list is not free — it is dictated by the metadata, in a fixed order. Getting it wrong gives you a link error at best and silent garbage at worst, so PSyclone will print the interface the metadata implies:
module mini_matvec_mod
implicit none
public
contains
subroutine mini_matvec_code(nlayers, field_1_w1, field_2_w3, ndf_w1, undf_w1, map_w1, ndf_w3, undf_w3, map_w3)
use constants_mod
integer(kind=i_def), intent(in) :: nlayers
integer(kind=i_def), intent(in) :: ndf_w1
integer(kind=i_def), dimension(ndf_w1), intent(in) :: map_w1
integer(kind=i_def), intent(in) :: ndf_w3
integer(kind=i_def), dimension(ndf_w3), intent(in) :: map_w3
integer(kind=i_def), intent(in) :: undf_w1
integer(kind=i_def), intent(in) :: undf_w3
real(kind=r_def), dimension(undf_w1), intent(inout) :: field_1_w1
real(kind=r_def), dimension(undf_w3), intent(in) :: field_2_w3
end subroutine mini_matvec_code
end module mini_matvec_modCompare with the hand-written subroutine in the same file: same arguments, same order, same shapes. This is the first thing to run when a new kernel does not work, and the fastest way to learn the argument-order convention — nlayers, then field arrays in meta_args order, then per-function-space ndf, undf, map.
The .x90 is a DSL, and violating its rules gets you an error at build time rather than a wrong answer at run time. Two you will hit:
bad = """
module bad_alg_mod
use constants_mod, only : r_def
use field_mod, only : field_type
implicit none
contains
subroutine bad_alg(lhs)
type(field_type), intent(inout) :: lhs
call invoke( inc_X_times_Y(lhs, lhs) )
end subroutine bad_alg
end module bad_alg_mod
"""
open("/tmp/bad_alg_mod.x90", "w").write(bad)
run("psyclone", "-api", "lfric", "--config", LFRIC_CONFIG,
"-oalg", "/tmp/bad_alg.f90", "-opsy", "/tmp/bad_psy.f90",
"/tmp/bad_alg_mod.x90", expect_failure=True)Generation Error: Argument 'lhs' is passed into kernel 'inc_x_times_y' code more than once from the algorithm layer. This is not allowed."Generation Error: Argument 'lhs' is passed into kernel 'inc_x_times_y' code more than once from the algorithm layer. This is not allowed.\n"The same field twice in one builtin — refused, because the generated loop would alias. The other rules worth knowing, from LFRic’s coding standards rather than from PSyclone’s checks:
invoke must be declared locally. Not fetched inline from a module, not returned from a function call in the argument list. PSyclone has to find the declaration to know the type, and it does not chase across files.invoke where you can. PSyclone can fuse loops and elide halo exchanges within an invoke and cannot see across two. demo/algorithm/mini_alg_mod.x90 has two on purpose; joining them would let the setval_c and the kernel loop be considered together.LFRic did not start from nothing. The atmosphere model inherits a large body of physics from the Unified Model — gravity-wave drag, boundary layer, microphysics, radiation — written years before PSyKAl existed, in ordinary Fortran, with ordinary loops over i, j and k. Rewriting it all as kernels is a decade of work that nobody has funded.
So lfric_apps runs PSyclone over that source in transformation mode: no invoke, no metadata, no generated PSy layer. Just “here is a routine, here is a script, add the OpenMP”. The name for this in the build system is transmute.
The wiring is worth knowing, because it tells you where to put a script:
| piece | what it does |
|---|---|
PSYCLONE_PHYSICS_FILES in applications/*/build/psyclone_transmute_file_list.mk |
the opt-in list, by module name. A file not named here is compiled untouched |
.xu90 |
the preprocessed intermediate; the physics is .F90 with real #ifs, and PSyclone will not run cpp for you |
interfaces/build/psyclone_transmute.mk |
the rules. Looks for <optimisation>/transmute/<subdir>/<module>.py, then local.py, then global.py, then falls back to a bare pass |
interfaces/build/transmute_psytrans/ |
the shared helper library, transmute_functions.py |
Note that last fallback: a file on the list with no script still goes through PSyclone, and comes out reformatted with its comments deleted and its existing OpenMP stripped ("Psyclone pass with no optimisation applied, OMP removed"). That is intentional — it guarantees that hand-written directives cannot survive into a build where the script is meant to own parallelism.
Ticket #813 applies PSyclone-driven OpenMP to gw_ussp_mod.F90, the Ultra-Simple Spectral Parametrization of gravity-wave drag. It is the best worked example available of what a performance ticket looks like end to end, so it is worth reproducing rather than summarising. The Exeter NG-ARCH wiki has the analysis behind it.
The script is 124 lines and the whole of its judgement is in three constants:
# -----------------------------------------------------------------------------
# (C) Crown copyright Met Office. All rights reserved.
# The file LICENCE, distributed with this code, contains details of the terms
# under which the code may be used.
# -----------------------------------------------------------------------------
"""
PSyclone Script for applying OpenMP transformations to Spectral Gravity Wave.
Provides:
- `trans(psyir)`: orchestrates all OpenMP transformations.
* Forces a dynamic-scheduled OMP **PARALLEL DO** over
`do i = 1, meta_segments%num_segments` with an explicit PRIVATE list.
* Clusters adjacent top-level loops into **PARALLEL regions** to amortize
thread overheads.
* Applies **STATIC-scheduled** OMP DO / PARALLEL DO to identified heavy
k-loops (profiling hotspots).
Context (specific to gw_ussp_mod.F90):
- Default scheduling is **static** everywhere **except** for the
`meta_segments%num_segments` loop, which is explicitly **dynamic** to
accommodate variable per-segment workloads.
- Certain scalars are initialised prior to OpenMP transformation so that
backends that emit FIRSTPRIVATE have a well-defined initial value.
Rationale:
- Profiling indicates k-index loops writing to heavy variables dominate
runtime. Static scheduling generally performs best for these loops, while
the meta-segmentation loop benefits from dynamic scheduling due to load
imbalance.
- Grouping adjacent outer loops into a single parallel region reduces
enter/exit overhead and improves locality.
Future Work:
- As PSyclone evolves, revisit nowait/lastprivate usage and refine region
boundaries.
"""
import logging
from psyclone.psyir.nodes import Routine
from transmute_psytrans.transmute_functions import (
add_parallel_do_over_meta_segments,
parallel_regions_for_clustered_loops,
omp_do_for_heavy_loops,
get_compiler,
)
# ------------------------------------------------------------------------------
# OpenMP transformation objects
#
# Policy:
# - STATIC schedule by default for heavy k-loops (best throughput observed).
# - DYNAMIC schedule **only** for the PARALLEL DO
# over meta_segments%num_segments.
# ------------------------------------------------------------------------------
# NOTE: OMP_* transform instances now live in the generic module; kept comment
# here to preserve file style and intent.
# Variables considered "heavy" when written inside loops:
# - HEAVY_VARS_K → for k-loops
# - HEAVY_VARS_I → for i-loops
# These guide where OMP DO / PARALLEL DO (static) should be applied.
HEAVY_VARS_K = {
"rhont_smallhalo", "rho_th", "udotk", "g_x", "g_y", "t_inc",
"r_u", "r_v", "g_xp_smallhalo", "g_yp_smallhalo",
}
HEAVY_VARS_I = {"rho_th", "nbv", "udotk"}
# --- meta_segments PARALLEL DO (forced, dynamic), with explicit privates ----
# Explicit PRIVATE list for the meta-segmentation PARALLEL DO
_META_PRIVATES = [
# scalars
"ii", "jj", "ss", "i", "k", "jdir",
# allocatables used per-segment
"s_sin_theta_lat", "s_totalppn", "s_rho_th", "s_nbv", "s_udotk", "s_fptot",
]
# ------------------------------------------------------------------------------
# Entry point
# ------------------------------------------------------------------------------
def trans(psyir):
"""
Entry point for OpenMP transformations for `gw_ussp_mod.F90`.
Passes:
1) Force PARALLEL DO (dynamic) for meta-segmentation loop.
2) Cluster adjacent top-level loops into PARALLEL regions.
3) Apply OMP DO / PARALLEL DO (static) to heavy k-loops.
"""
logging.info(
"gw_ussp_mod.F90 Psyclone Optimisation for the CPU is running."
)
compiler = get_compiler()
# GCC currently has issues with firstprivated indexes,
# which is soon to be resolved in PSyclone.
# However for now we will avoid using OpenMP
# around meta_segments with with GCC.
if compiler == "gnu":
logging.info(
"Skipping gw_ussp_mod meta_segment optimisations for GCC."
)
else:
# 1) Force a PARALLEL DO around meta_segments%num_segments (dynamic)
for routine in psyir.walk(Routine):
add_parallel_do_over_meta_segments(
routine,
container_name="meta_segments",
member_name="num_segments",
privates=_META_PRIVATES,
)
# 2) Cluster adjacent outer loops into PARALLEL regions (no schedule here)
for routine in psyir.walk(Routine):
parallel_regions_for_clustered_loops(routine)
# 3) Add OMP DO / PARALLEL DO on heavy k- and i-loops (static)
# (Skip the special i-loop handled above.)
for routine in psyir.walk(Routine):
omp_do_for_heavy_loops(routine, "k", HEAVY_VARS_K)
omp_do_for_heavy_loops(
routine, "i", HEAVY_VARS_I,
skip_member_count=("i", "meta_segments", "num_segments")
)HEAVY_VARS_K and HEAVY_VARS_I are lists of variable names. A loop is worth parallelising if it writes to one of them. That is the ticket’s whole model of the routine, and it came from a profiler: these are the variables the hot loops write.
The experiment to run. Comparing the transformed output against the original source is useless — the diff drowns in reformatting and deleted comments. Compare it instead against PSyclone’s output with no script, which isolates exactly what the script did:
WARNING:root:OMP applied to k-loop (static).
WARNING:root:OMP applied to k-loop (static).
WARNING:root:OMP applied to k-loop (static).
WARNING:root:OMP applied to k-loop (static).
WARNING:root:OMP applied to k-loop (static).
WARNING:root:OMP applied to k-loop (static).
WARNING:root:OMP applied to k-loop (static).
WARNING:root:OMP applied to k-loop (static).
WARNING:root:OMP applied to i-loop (static).
WARNING:root:OMP applied to i-loop (static).
WARNING:root:OMP applied to i-loop (static).'WARNING:root:OMP applied to k-loop (static).\nWARNING:root:OMP applied to k-loop (static).\nWARNING:root:OMP applied to k-loop (static).\nWARNING:root:OMP applied to k-loop (static).\nWARNING:root:OMP applied to k-loop (static).\nWARNING:root:OMP applied to k-loop (static).\nWARNING:root:OMP applied to k-loop (static).\nWARNING:root:OMP applied to k-loop (static).\nWARNING:root:OMP applied to i-loop (static).\nWARNING:root:OMP applied to i-loop (static).\nWARNING:root:OMP applied to i-loop (static).\n'--- psyclone, no script
+++ psyclone, ticket #813 script
@@ -158,4 +158,6 @@
call v_to_p(v, vdims_s%i_start, vdims_s%i_end, vdims_s%j_start, vdims_s%j_end, pdims%i_start, pdims%i_end, pdims%j_start, &
&pdims%j_end, levels, vonp)
+ !$omp parallel default(shared) private(i,j,jdir,k)
+ !$omp do schedule(static)
do k = pdims_s%k_start, pdims_s%k_end, 1
do j = pdims_s%j_start, pdims_s%j_end, 1
@@ -166,4 +168,6 @@
enddo
enddo
+ !$omp end do
+ !$omp do schedule(static)
do k = pdims%k_start, pdims%k_end, 1
do j = pdims%j_start, pdims%j_end, 1
@@ -174,4 +178,6 @@
enddo
enddo
+ !$omp end do
+ !$omp do schedule(static)
do k = tkfix1start + 1, tdims%k_end - 1, 1
do j = tdims%j_start, tdims%j_end, 1
@@ -185,4 +191,6 @@
enddo
enddo
+ !$omp end do
+ !$omp do schedule(static)
do i = tdims%i_start, tdims%i_end, 1
j = 1
@@ -192,5 +200,7 @@
nbv(i,j,tdims%k_end) = nbv(i,j,tdims%k_end - 1)
enddo
+ !$omp end do
do k = tdims%k_end, tkfix1start + 1, -1
+ !$omp do schedule(static)
do i = tdims%i_start, tdims%i_end, 1
j = 1
@@ -199,6 +209,8 @@
end if
enddo
+ !$omp end do
enddo
do jdir = 1, idir, 1
+ !$omp do schedule(static)
do k = tkfix1start, tdims%k_end - 1, 1
do j = tdims%j_start, tdims%j_end, 1
@@ -209,9 +221,13 @@
enddo
enddo
+ !$omp end do
+ !$omp do schedule(static)
do i = tdims%i_start, tdims%i_end, 1
j = 1
udotk(i,j,tdims%k_end,jdir) = uonp(i,j,pdims%k_end) * cosphi(jdir) + vonp(i,j,pdims%k_end) * sinphi(jdir)
enddo
- enddo
+ !$omp end do
+ enddo
+ !$omp end parallel
minlaunchlev = launchlev
maxlaunchlev = launchlev
@@ -222,4 +238,8 @@
ALLOCATE(segments(1:meta_segments%num_segments))
call segments_mod_segments(segments, meta_segments, row_length)
+ jdir = 0
+ k = 0
+ !$omp parallel do default(shared) private(i,ii,jj,s_fptot,s_nbv,s_rho_th,s_sin_theta_lat,s_totalppn,s_udotk,ss) &
+!$omp& firstprivate(jdir,k) schedule(dynamic)
do i = 1, meta_segments%num_segments, 1
ALLOCATE(s_sin_theta_lat(1:segments(i)%seg_points))
@@ -302,5 +322,8 @@
DEALLOCATE(s_sin_theta_lat)
enddo
+ !$omp end parallel do
DEALLOCATE(segments)
+ !$omp parallel default(shared) private(g_g,i,j,jdir,k)
+ !$omp do schedule(static)
do k = minlaunchlev + 1, pdims%k_end, 1
do j = pdims%j_start, pdims%j_end, 1
@@ -314,4 +337,6 @@
enddo
enddo
+ !$omp end do
+ !$omp do schedule(static)
do k = minlaunchlev, pdims%k_end, 1
do j = pdims%j_start, pdims%j_end, 1
@@ -322,5 +347,8 @@
enddo
enddo
+ !$omp end do
+ !$omp end parallel
if (l_ussp_heating) then
+ !$omp parallel do default(shared) private(dzb,dzl,dzu,i,j,k,ududt,uhat,vdvdt,vhat) schedule(static)
do k = tkfix1start, tdims%k_end - 1, 1
do j = tdims%j_start, tdims%j_end, 1
@@ -337,5 +365,7 @@
enddo
enddo
- end if
+ !$omp end parallel do
+ end if
+ !$omp parallel do default(shared) private(i,j,k) schedule(static)
do k = minlaunchlev, udims%k_end, 1
do j = udims%j_start, udims%j_end, 1
@@ -350,4 +380,5 @@
enddo
enddo
+ !$omp end parallel do
if (gwspec_nflux_on) then
do k = tkfix1start, tdims%k_end, 1Everything the ticket did, in about 110 lines of diff. Four things in it repay close reading.
One parallel region, many worksharing loops. The first hunk opens !$omp parallel and then issues a series of !$omp do / !$omp end do pairs across adjacent loops, closing the region only at the end. That is parallel_regions_for_clustered_loops, and the point is amortisation: entering a parallel region costs a thread-team synchronisation, and doing it once for eight loops instead of eight times is most of the win on short loops. The clustering is purely positional — “adjacent top-level loops under the same parent” — which is why the helper is 90 lines of index bookkeeping.
Worksharing inside a sequential loop. Look at this shape:
The k loop is not parallel — it walks downward and each level reads the level above, a loop-carried dependence. So the k loop is executed redundantly by every thread, and only the i loop is shared out. This is correct, and it is correct because of the implicit barrier at the end of !$omp do: every thread finishes level k before any starts k-1. It is also the kind of code nobody writes by hand, and a good demonstration that the machine-generated version can be better than the hand-rolled one it replaced.
firstprivate, and a compiler-specific branch. The second hunk is the interesting one:
Three things at once. The schedule is dynamic, alone in this file, because segments contain different numbers of launch levels and the work per iteration varies — classic load imbalance, and the one place where dynamic scheduling earns its overhead. The private list is explicit, supplied by the script, because it includes per-segment allocatables that PSyclone’s analysis would not privatise on its own. And jdir = 0; k = 0 were inserted before the region — the script initialises them so that the firstprivate PSyclone emits has a defined value to copy.
That firstprivate is also why the script branches on the compiler:
run("psyclone", "-l", "all", "-s", ticket_script, "-o", "/tmp/gw_gnu.f90", gw_source,
env={"PYTHONPATH": str(LFRIC_TRANSMUTE_TOOLS), "COMPILER": "gnu"})
print("dynamic PARALLEL DO present under CCE:",
"schedule(dynamic)" in open("/tmp/gw_ticket.f90").read())
print("dynamic PARALLEL DO present under GNU:",
"schedule(dynamic)" in open("/tmp/gw_gnu.f90").read())WARNING:root:OMP applied to k-loop (static).
WARNING:root:OMP applied to k-loop (static).
WARNING:root:OMP applied to k-loop (static).
WARNING:root:OMP applied to k-loop (static).
WARNING:root:OMP applied to k-loop (static).
WARNING:root:OMP applied to k-loop (static).
WARNING:root:OMP applied to k-loop (static).
WARNING:root:OMP applied to k-loop (static).
WARNING:root:OMP applied to i-loop (static).
WARNING:root:OMP applied to i-loop (static).
WARNING:root:OMP applied to i-loop (static).
dynamic PARALLEL DO present under CCE: True
dynamic PARALLEL DO present under GNU: FalseUnder GNU the whole meta_segments pass is skipped — “GCC currently has issues with firstprivated indexes”, says the script — so the same source, the same PSyclone and the same script produce differently parallelised code for different compilers. get_compiler() reads FC, CC or the loaded module list to decide. This is the thing to hold on to about optimisation scripts: they are not a description of the program, they are a description of the program on a machine, and the machine is a parameter.
What is not there. No collapse(). The ticket says so explicitly: collapse helped GCC and hurt CCE, and with only one script shared between them, the loss on the production compiler decided it. The ticket also deleted the routine’s hand-rolled omp_block machinery, which had been 0.5% slower than nothing at all. Both are worth noting as the normal outcome of this work: a ticket’s value is often a removal plus a decision not to use a feature.
transmute_functions.py is the general-purpose part, and it is a good model for how these scripts are actually built — a predicate over the tree, then a transformation applied where the predicate holds:
# -----------------------------------------------------------------------------
# (C) Crown copyright Met Office. All rights reserved.
# The file LICENCE, distributed with this code, contains details of the terms
# under which the code may be used.
# -----------------------------------------------------------------------------
'''
Transmute functions for optimisation scripts. Intension is for these to
eventually live in the PSyTran repo. See Ticket #906.
'''
import logging
import os
from itertools import dropwhile, takewhile
from typing import Sequence, Optional, Tuple, Set
from psyclone.psyir.nodes import (
Node,
Loop,
Call,
Assignment,
Reference,
OMPParallelDirective,
OMPDoDirective,
OMPParallelDoDirective,
StructureReference,
Member,
Literal,
Schedule,
)
from psyclone.psyir.symbols import (
DataSymbol,
ContainerSymbol,
RoutineSymbol,
ImportInterface,
UnsupportedFortranType,
ScalarType
)
from psyclone.psyir.transformations import OMPParallelTrans
from psyclone.transformations import (
OMPLoopTrans,
TransformationError,
OMPParallelLoopTrans,
)
# ------------------------------------------------------------------------------
# OpenMP transformation objects
#
# Policy:
# - STATIC schedule by default for heavy loops (best throughput observed).
# - DYNAMIC schedule **only** for the special PARALLEL DO case that the app
# identifies (e.g. a member-count loop like meta_segments%num_segments).
# ------------------------------------------------------------------------------
OMP_PARALLEL_REGION_TRANS = OMPParallelTrans()
# Default: static schedule for heavy loops
OMP_DO_LOOP_TRANS_STATIC = OMPLoopTrans(omp_schedule="static")
OMP_PARALLEL_LOOP_DO_TRANS_STATIC = OMPParallelLoopTrans(
omp_schedule="static", omp_directive="paralleldo"
)
# Exception: dynamic schedule for the app-selected special loop
OMP_PARALLEL_LOOP_DO_TRANS_DYNAMIC = OMPParallelLoopTrans(
omp_schedule="dynamic", omp_directive="paralleldo"
)
OMP_DO_LOOP_TRANS_DYNAMIC = OMPLoopTrans(omp_schedule="dynamic")
# ------------------------------------------------------------------------------
def is_heavy_loop(loop, heavy_vars: Set[str]):
"""
Determine whether `loop` performs significant work on key variables.
Parameters
----------
loop : psyclone.psyir.nodes.Loop
Candidate loop to inspect.
heavy_vars : set[str]
Names of variables (profiling hotspots) that mark a loop as heavy.
Returns
-------
bool
True if any Assignment within the loop writes to a Reference whose
name is in `heavy_vars`; False otherwise.
"""
for assign in loop.walk(Assignment):
lhs = assign.lhs
if isinstance(lhs, Reference) and lhs.name in heavy_vars:
return True
return Falseis_heavy_loop is eight lines and is the entire targeting mechanism: walk the loop’s Assignment nodes, look at the left-hand side, ask whether its name is in a set. Crude — it cannot tell a hot loop from a cold one, only a loop that touches a variable a human already identified — and exactly the right amount of machinery, because the profiler did the analysis and the script only has to encode the conclusion.
These helpers are being upstreamed into PSyTran (issue #118), with ticket #906 tracking deletion of the local copies. If you write generic tree predicates, that is where they should end up.
The one part of PSyclone that is not about parallelism. PSyAD reads a tangent-linear kernel and writes its adjoint.
The mathematics first, since it is the shortest route in. A nonlinear model step is M: \mathbb{R}^n \to \mathbb{R}^n. Linearising about a state \bar{x} gives the tangent-linear operator L = \partial M/\partial x|_{\bar{x}}, a matrix acting on perturbations: \delta y = L\,\delta x. Its adjoint is L^{*}, defined by
\langle L\,\delta x, \delta y\rangle = \langle \delta x, L^{*}\,\delta y\rangle \quad \forall\, \delta x, \delta y.
In real arithmetic L^{*} = L^{\mathsf{T}}. You know both of these as jax.jvp and jax.vjp: the tangent-linear is forward mode, the adjoint is reverse mode, and the identity above is the statement that makes backpropagation work. Variational data assimilation needs L^{*} for the same reason training needs a gradient — 4D-Var minimises a cost function over the initial state, and the gradient of that cost comes back through the adjoint of the whole forecast.
The difference from JAX is when. JAX differentiates a traced computation at run time, in Python. PSyAD reads Fortran and writes Fortran, ahead of time, and what gets compiled into the model is an ordinary subroutine that a Fortran programmer can read and profile.
tl_kernel = """
module tl_mix_mod
implicit none
contains
subroutine tl_mix_code(nlayers, ls_theta, theta, rho, dz)
integer, intent(in) :: nlayers
real, intent(in) :: ls_theta(nlayers)
real, intent(inout) :: theta(nlayers)
real, intent(in) :: rho(nlayers)
real, intent(in) :: dz
integer :: k
do k = 2, nlayers
theta(k) = theta(k) + dz * ls_theta(k) * (rho(k) - rho(k-1))
end do
end subroutine tl_mix_code
end module tl_mix_mod
"""
open("/tmp/tl_mix_mod.f90", "w").write(tl_kernel)
run("psyad", "-a", "theta", "rho",
"-oad", "/tmp/adj_mix_mod.f90", "-otest", "/tmp/adj_test.f90",
"/tmp/tl_mix_mod.f90")
show("/tmp/adj_mix_mod.f90")module adj_mix_mod
implicit none
public
contains
subroutine adj_mix_code(nlayers, ls_theta, theta, rho, dz)
integer, intent(in) :: nlayers
real, dimension(nlayers), intent(in) :: ls_theta
real, dimension(nlayers), intent(in) :: theta
real, dimension(nlayers), intent(inout) :: rho
real, intent(in) :: dz
integer :: k
do k = nlayers, 2, -1
rho(k) = rho(k) + dz * ls_theta(k) * theta(k)
rho(k - 1) = rho(k - 1) - dz * ls_theta(k) * theta(k)
enddo
end subroutine adj_mix_code
end module adj_mix_modRead the reversal, because it is the whole algorithm in four lines:
| tangent-linear | adjoint |
|---|---|
do k = 2, nlayers |
do k = nlayers, 2, -1 |
theta is intent(inout) |
theta is intent(in) |
rho is intent(in) |
rho is intent(inout) |
theta(k) += dz*ls_theta(k)*(rho(k) - rho(k-1)) |
rho(k) += dz*ls_theta(k)*theta(k); rho(k-1) -= dz*ls_theta(k)*theta(k) |
The loop runs backwards; the roles of input and output swap; and each read in the forward code becomes an increment in the reverse code, with the sign it had. Transposition of a matrix, expressed as a statement reordering.
Note ls_theta — the “linearisation state” — is untouched: it is a passive variable, part of \bar{x} rather than of \delta x. -a theta rho declared the active ones, and that list is the one thing PSyAD cannot infer for you. Getting it wrong is the classic PSyAD error: name too few and the adjoint is incomplete, too many and you get an adjoint of something you did not mean.
-otest generated a harness alongside the adjoint. It applies L to random \delta x, applies L^{*} to the result, and checks the inner-product identity numerically. That is the standard adjoint correctness test, and it is a genuinely strong one — an adjoint with a sign error or a missing term fails it immediately.
program adj_test
use tl_mix_mod, only : tl_mix_code
use adj_mix_mod, only : adj_mix_code
integer, parameter :: array_extent = 20
integer, parameter :: nlayers = array_extent
real, parameter :: overall_tolerance = 1500.0
real :: inner1
real :: inner2
real, dimension(nlayers) :: ls_theta
real, dimension(nlayers) :: ls_theta_input
real, dimension(nlayers) :: theta
real, dimension(nlayers) :: theta_input
real, dimension(nlayers) :: rho
real, dimension(nlayers) :: rho_input
real :: dz
real :: dz_input
real :: MachineTol
real :: relative_diff
! Initialise the kernel arguments and keep copies of them
call random_number(ls_theta)
ls_theta_input = ls_theta
call random_number(theta)
theta_input = theta
call random_number(rho)
rho_input = rho
call random_number(dz)
dz_input = dz
! Call the tangent-linear kernel
call tl_mix_code(nlayers, ls_theta, theta, rho, dz)
! Compute the inner product of the results of the tangent-linear kernel
inner1 = 0.0
inner1 = inner1 + DOT_PRODUCT(ls_theta, ls_theta)
inner1 = inner1 + DOT_PRODUCT(theta, theta)
inner1 = inner1 + DOT_PRODUCT(rho, rho)
inner1 = inner1 + dz * dz
! Call the adjoint of the kernel
call adj_mix_code(nlayers, ls_theta, theta, rho, dz)
! Compute inner product of results of adjoint kernel with the original inputs to the tangent-linear kernel
inner2 = 0.0
inner2 = inner2 + DOT_PRODUCT(ls_theta, ls_theta_input)
inner2 = inner2 + DOT_PRODUCT(theta, theta_input)
inner2 = inner2 + DOT_PRODUCT(rho, rho_input)
inner2 = inner2 + dz * dz_input
! Test the inner-product values for equality, allowing for the precision of the active variables
MachineTol = SPACING(MAX(ABS(inner1), ABS(inner2)))
relative_diff = ABS(inner1 - inner2) / MachineTol
if (relative_diff < overall_tolerance) then
! PSyclone CodeBlock (unsupported code) reason:
! - Unsupported statement: Write_Stmt
WRITE(*, *) 'Test of adjoint of ''tl_mix_code'' PASSED: ', inner1, inner2, relative_diff
else
! PSyclone CodeBlock (unsupported code) reason:
! - Unsupported statement: Write_Stmt
WRITE(*, *) 'Test of adjoint of ''tl_mix_code'' FAILED: ', inner1, inner2, relative_diff
end if
end program adj_test
Test of adjoint of 'tl_mix_code' PASSED: 20.3455353 20.3455353 0.00000000" Test of adjoint of 'tl_mix_code' PASSED: 20.3455353 20.3455353 0.00000000 \n"Two inner products agreeing to single-precision tolerance. The whole chain — generate the adjoint, generate its test, compile, run — is four commands, and lfric_apps runs the LFRic-flavoured version of it in science/adjoint as part of its test suite.
The failure modes have a short tail, and knowing which layer failed tells you what to do.
| symptom | layer | usual cause |
|---|---|---|
TransformationError: ... cannot be enclosed by ... |
transformation | a CodeBlock in the region — print them |
TransformationError: ... are dependent and cannot be parallelised |
dependence analysis | a real dependence, or an indirection it cannot see through |
| transformation silently applied to nothing | your script | the predicate never matched; walk() found no loop of that name |
ImportError on a *Trans |
version skew | see §3.3; names moved in 3.3 |
Generation Error: ... in PSyKAl mode |
the DSL’s own rules | metadata/algorithm mismatch, or a rule like the double-argument one in §4.6 |
| unresolved symbol in generated code | frontend | a use PSyclone could not follow; -d and -I tell it where to look |
| wrong loop bounds in a PSy layer | config, not script | --config, and COMPUTE_ANNEXED_DOFS in it |
Three tools for the middle column. tree(node) — the tree as PSyclone sees it, which is the ground truth about what your predicate is matching against. node.debug_string() — the Fortran for one node, without needing a whole valid file, and it works on trees mid-transformation that the backend would refuse to write. And a plain breakpoint() in your trans function: it is imported Python, so the debugger works normally, and stepping through a script on a real file is the fastest way to learn the API.
do i = 1, n, 1
a(i) = a(i) * 2.0
enddo
position: 0 | depth: 3 | parent: Routine
ancestor Routine: two_passesposition, depth, parent, ancestor, walk and detach are the six methods most of a script is made of.
No answers. Each is a few lines in a cell above, or a psyclone invocation.
diffuse routine from §2 and parallelise the k loop. Then try the i loop instead, and explain the refusal in terms of the stencil.write(*,*) statement to the middle of the diffuse loop nest and find out which transformations you have just lost. Then move it out and check you have them back.LoopSwapTrans swaps a loop nest. Apply it to diffuse and decide, from the Fortran note’s column-major discussion, which order you want and by how much.demo/algorithm/mini_alg_mod.x90, join the two invokes into one. Diff the generated PSy layer against the two-invoke version. What did PSyclone gain by seeing them together, and did the halo exchange move?GH_INC to GH_READWRITE in the demo kernel and regenerate. The colouring should disappear — and so should something else in the loop bounds. Why is GH_READWRITE on W1 a lie?trans that finds every loop in gw_ussp_mod.F90 writing to a variable of your choice, and reports the loop variable and nesting depth without transforming anything. This is the reconnaissance half of a real ticket.--keep-comments. Do the comments end up attached to the right statements after the OpenMP is inserted?psyad a tangent-linear kernel with a where construct, or a division by an active variable, and see what it does. Then check whether the generated harness still passes.psyclone -p routines inserts profiling hooks. Add it to the demo and read what the instrumentation costs in generated code.lfric_core, a kernel whose operates_on is not CELL_COLUMN. Generate its stub and account for every argument.| you know | PSyclone | the difference that matters |
|---|---|---|
| jaxpr | PSyIR | statement-shaped, because it must print back out as reviewable Fortran |
objmode |
CodeBlock |
opaque, and it silently blocks transformations of any region containing it |
vmap/pmap |
the generated PSy layer | ahead of time, from metadata, into a file you can read |
jax.vjp |
PSyAD | source-to-source, with a generated correctness test |
-O2 |
nothing | there is no optimiser; a script names every rewrite |
| a compiler flag | optimisation/<site>/<dsl>/*.py |
reviewed, versioned, and per-site |
| a race condition | TransformationError |
refusal is the product |
PSyKAl (-api lfric): metadata plus an invoke is enough for PSyclone to write the loops, halo exchanges and dirty flags, and your script decides how they are parallelised.
Transmute (no -api): PSyclone rewrites the loops you already have, and your script has to identify them itself — usually from a profile.
CodeBlock you did not know was there. Transformations refuse, or worse, apply to a smaller region than you meant. Print them first.psyclone.transformations versus psyclone.psyir.transformations, and Dynamo0p3* versus LFRic*.--config in PSyKAl mode, and getting loop bounds derived from PSyclone’s defaults instead of LFRic’s.--keep-comments.submodules/PSyclone/tutorial/ — the official hands-on tutorial, with notebooks for the fparser2 and PSyIR layers and practicals for both the generic and LFRic routes. Complementary to this note: it teaches the tool on its own terms, at more length.submodules/PSyclone/examples/lfric/ — twenty numbered examples, each a minimal .x90 plus a script, covering builtins, stencils, operators, kernel transformations and GPU offload.lfric_core/infrastructure/build/psyclone/psyclone_tools.py and the optimisation/ tree of any lfric_apps application — the production scripts, and the best guide to house style.The alternative — pragmas in the science code — was tried by everyone and is what LFRic is deliberately escaping. A !$omp parallel do collapse(2) schedule(dynamic) in a physics routine is a statement about one machine, welded to code that has to run on five. Ticket #813, the case study in §6, is partly the story of removing such hand-rolled machinery (omp_block) and replacing it with a script.↩︎
Two places in the submodules still disagree with reality, both worth a patch. lfric_core/documentation/source/getting_started/installation/software_dependencies.rst still says “PSyclone (3.2.2)” and “fparser (0.2.1)” — the latter is impossible, since psyclone 3.3.1 requires exactly fparser==0.2.4. And lfric_apps/rose-stem/site/azngarch/common/suite_config_sandbox.cylc still loads py-psyclone/3.1.0.↩︎
This is the standard trade-off in source-to-source tools and the reason they are a distinct genre. ROSE, Coccinelle and clang’s AST-matcher refactorings sit in the same place; LLVM IR and jaxpr sit on the other side of the line. The tell is whether the IR can be printed back out as something a maintainer would accept in a pull request.↩︎
With the caveat that a comment’s position after a transformation is guesswork — a comment attached to a loop that has been fused, colored or wrapped in a directive can end up describing the wrong thing. That is presumably why the flag is opt-in.↩︎
“Dynamo” was the name of the project that became LFRic’s dynamical core, and dynamo0.3 was the API name for a decade — long enough that it is still what much of the PSyclone test suite and documentation calls it. If you grep for it, that is why.↩︎
LFRic can also tile the colouring (options={"tiling": True} on LFRicColourTrans), which recovers some locality by colouring blocks of cells rather than single cells. psyclone_tools.colour_loops takes an enable_tiling flag for exactly this.↩︎