def trans(psy):
''' PSyclone transformation script for the dynamo0p3 API to apply
loop fusion and OpenMP for a particular example.'''
otrans = DynamoOMPParallelLoopTrans()
ftrans = DynamoLoopFuseTrans()
invoke = psy.invokes.invoke_list[0]
schedule = invoke.schedule
from psyclone.configuration import Config
config = Config.get()
if config.api_conf("dynamo0.3").compute_annexed_dofs and \
config.distributed_memory:
# We can't loop fuse as the loop bounds differ so add
# OpenMP parallel do directives to the loops
schedule, _ = otrans.apply(schedule.children[0])
schedule, _ = otrans.apply(schedule.children[1])
else:
# Loop fuse the two built-in kernels. The 'same_space' flag needs to
# be set as built-ins are over ANY_SPACE.
ftrans.same_space = True
schedule, _ = ftrans.apply(schedule[0], schedule[1])
# Add an OpenMP parallel do directive to the resultant loop-fused loop
schedule, _ = otrans.apply(schedule.children[0])
# take a look at what we've done
schedule.view()
schedule.dag(file_format="png")
return psy
def trans(psy):
''' PSyclone transformation script for the dynamo0p3 api to apply
colouring and OpenMP generically.'''
ctrans = Dynamo0p3ColourTrans()
otrans = DynamoOMPParallelLoopTrans()
# Loop over all of the Invokes in the PSy object
for invoke in psy.invokes.invoke_list:
print("Transforming invoke '" + invoke.name + "'...")
schedule = invoke.schedule
# Colour all of the loops over cells unless they are on
# discontinuous spaces (W3, WTHETA and W2V)
cschedule = schedule
for child in schedule.children:
if isinstance(child, Loop) \
and child.field_space.orig_name \
not in DISCONTINUOUS_FUNCTION_SPACES \
and child.iteration_space == "cells":
cschedule, _ = ctrans.apply(child)
# Then apply OpenMP to each of the colour loops
schedule = cschedule
for child in schedule.children:
if isinstance(child, Loop):
if child.loop_type == "colours":
schedule, _ = otrans.apply(child.children[0])
else:
schedule, _ = otrans.apply(child)
schedule.view()
invoke.schedule = schedule
return psy
def test_no_colours_loop_dynamo0p3():
''' Test that applying LFRicExtractTrans on a Loop over cells
in a colour without its parent Loop over colours in Dynamo0.3 API
raises a TransformationError. '''
etrans = LFRicExtractTrans()
ctrans = Dynamo0p3ColourTrans()
otrans = DynamoOMPParallelLoopTrans()
_, invoke = get_invoke("1_single_invoke.f90",
DYNAMO_API,
idx=0,
dist_mem=False)
schedule = invoke.schedule
# Colour first loop that calls testkern_code (loop is over cells and
# is not on a discontinuous space)
ctrans.apply(schedule.children[0])
colour_loop = schedule.children[0].loop_body[0]
# Apply OMP Parallel DO Directive to the colour Loop
otrans.apply(colour_loop)
directive = schedule[0].loop_body
# Try to extract the region between the Loop over cells in a colour
# and the exterior Loop over colours
with pytest.raises(TransformationError) as excinfo:
_, _ = etrans.apply(directive)
assert ("Dynamo0.3 API: Extraction of a Loop over cells in a "
"colour without its ancestor Loop over colours is not "
"allowed.") in str(excinfo.value)
def trans(psy):
'''PSyclone transformation script for the dynamo0p3 API that applies
OpenMP parallel loop parallelisation. It also outputs a textual
representation of the transformated PSyIR.
:param psy: a PSyclone PSy object which captures the algorithm and \
kernel information required by PSyclone.
:type psy: subclass of :py:class:`psyclone.psyGen.PSy`
'''
otrans = DynamoOMPParallelLoopTrans()
for invoke in psy.invokes.invoke_list:
schedule = invoke.schedule
# Add OpenMP parallel do directives to the loops
for loop in schedule.loops():
try:
otrans.apply(loop)
except TransformationError as info:
print(str(info.value))
# take a look at what we've done
schedule.view()
return psy
def test_no_colours_loop_dynamo0p3():
''' Test that applying DynamoExtractRegionTrans on a Loop over cells
in a colour without its parent Loop over colours in Dynamo0.3 API
raises a TransformationError. '''
from psyclone.transformations import Dynamo0p3ColourTrans, \
DynamoOMPParallelLoopTrans
etrans = DynamoExtractRegionTrans()
ctrans = Dynamo0p3ColourTrans()
otrans = DynamoOMPParallelLoopTrans()
_, invoke_info = parse(os.path.join(DYNAMO_BASE_PATH,
"1_single_invoke.f90"),
api=DYNAMO_API)
psy = PSyFactory(DYNAMO_API, distributed_memory=False).create(invoke_info)
invoke = psy.invokes.invoke_list[0]
schedule = invoke.schedule
# Colour first loop that calls testkern_code (loop is over cells and
# is not on a discontinuous space)
schedule, _ = ctrans.apply(schedule.children[0])
colour_loop = schedule.children[0].loop_body[0]
# Apply OMP Parallel DO Directive to the colour Loop
schedule, _ = otrans.apply(colour_loop)
directive = schedule.children[0].children[0]
# Try to extract the region between the Loop over cells in a colour
# and the exterior Loop over colours
with pytest.raises(TransformationError) as excinfo:
_, _ = etrans.apply(directive)
assert ("Dynamo0.3 API: Extraction of a Loop over cells in a "
"colour without its ancestor Loop over colours is not "
"allowed.") in str(excinfo)
def trans(psy):
''' PSyclone transformation script for the dynamo0p3 api to apply
colouring and OpenMP generically.'''
ctrans = Dynamo0p3ColourTrans()
otrans = DynamoOMPParallelLoopTrans()
# Loop over all of the Invokes in the PSy object
for invoke in psy.invokes.invoke_list:
print "Transforming invoke '" + invoke.name + "'..."
schedule = invoke.schedule
# Colour all of the loops over cells unless they are on W3
cschedule = schedule
for child in schedule.children:
if isinstance(child, Loop) and child.field_space != "w3" \
and child.iteration_space == "cells":
cschedule, _ = ctrans.apply(child)
# Then apply OpenMP to each of the colour loops
schedule = cschedule
for child in schedule.children:
if isinstance(child, Loop):
if child.loop_type == "colours":
schedule, _ = otrans.apply(child.children[0])
else:
schedule, _ = otrans.apply(child)
schedule.view()
invoke.schedule = schedule
return psy
def test_extract_single_builtin_dynamo0p3():
''' Test that extraction of a BuiltIn in an Invoke produces the
correct result in Dynamo0.3 API without and with optimisations. '''
from psyclone.transformations import DynamoOMPParallelLoopTrans
etrans = DynamoExtractRegionTrans()
otrans = DynamoOMPParallelLoopTrans()
# Test extract without optimisations
_, invoke_info = parse(os.path.join(
DYNAMO_BASE_PATH, "15.1.2_builtin_and_normal_kernel_invoke.f90"),
api=DYNAMO_API)
psy = PSyFactory(DYNAMO_API, distributed_memory=False).create(invoke_info)
invoke = psy.invokes.invoke_list[0]
schedule = invoke.schedule
schedule, _ = etrans.apply(schedule.children[1])
code = str(psy.gen)
output = (" ! ExtractStart\n"
" ! CALL write_extract_arguments(argument_list)\n"
" !\n"
" DO df=1,undf_any_space_1_f2\n"
" f2_proxy%data(df) = 0.0\n"
" END DO \n"
" !\n"
" ! ExtractEnd\n")
assert output in code
# Test extract with OMP Parallel optimisation
_, invoke_info = parse(os.path.join(
DYNAMO_BASE_PATH, "15.1.1_builtin_and_normal_kernel_invoke_2.f90"),
api=DYNAMO_API)
psy = PSyFactory(DYNAMO_API, distributed_memory=False).create(invoke_info)
invoke = psy.invokes.invoke_list[0]
schedule = invoke.schedule
schedule, _ = otrans.apply(schedule.children[1])
schedule, _ = etrans.apply(schedule.children[1])
code = str(psy.gen)
output = (" ! ExtractStart\n"
" ! CALL write_extract_arguments(argument_list)\n"
" !\n"
" !$omp parallel do default(shared), private(df), "
"schedule(static)\n"
" DO df=1,undf_any_space_1_f1\n"
" f1_proxy%data(df) = 0.5*f1_proxy%data(df) + "
"f2_proxy%data(df)\n"
" END DO \n"
" !$omp end parallel do\n"
" !\n"
" ! ExtractEnd\n")
assert output in code
def test_loop_no_directive_dynamo0p3():
''' Test that applying Extract Transformation on a Loop without its
parent Directive when optimisations are applied in Dynamo0.3 API
raises a TransformationError. '''
etrans = LFRicExtractTrans()
# Test a Loop nested within the OMP Parallel DO Directive
_, invoke = get_invoke("4.13_multikernel_invokes_w3_anyd.f90",
DYNAMO_API, idx=0, dist_mem=False)
schedule = invoke.schedule
# Apply DynamoOMPParallelLoopTrans to the second Loop
otrans = DynamoOMPParallelLoopTrans()
schedule, _ = otrans.apply(schedule[1])
loop = schedule.children[1].dir_body[0]
# Try extracting the Loop inside the OMP Parallel DO region
with pytest.raises(TransformationError) as excinfo:
_, _ = etrans.apply(loop)
assert "Error in LFRicExtractTrans: Application to a Loop without its " \
"parent Directive is not allowed." in str(excinfo.value)
def trans(psy):
'''PSyclone transformation script for the dynamo0p3 API that applies
loop colouring and OpenMP parallel loop parallelisation. It also
outputs a textual representation of the transformated PSyIR.
:param psy: a PSyclone PSy object which captures the algorithm and \
kernel information required by PSyclone.
:type psy: subclass of :py:class:`psyclone.psyGen.PSy`
'''
otrans = DynamoOMPParallelLoopTrans()
ctrans = Dynamo0p3ColourTrans()
ptrans = OMPParallelTrans()
ltrans = Dynamo0p3OMPLoopTrans()
const = LFRicConstants()
for invoke in psy.invokes.invoke_list:
schedule = invoke.schedule
# Colour any loops that need colouring
for loop in schedule.walk(Loop):
if (loop.field_space.orig_name
not in const.VALID_DISCONTINUOUS_NAMES
and loop.iteration_space == "cell_column"):
ctrans.apply(loop)
# Add OpenMP parallel do directives to the loops
for loop in schedule.walk(Loop):
try:
# Make sure reductions are reproducible
if loop.reductions():
ptrans.apply(loop)
ltrans.apply(loop, {"reprod": True})
else:
otrans.apply(loop)
except TransformationError as info:
print(str(info.value))
# take a look at what we've done
schedule.view()
return psy
def trans(psy):
''' PSyclone transformation script for the dynamo0p3 api to apply
colouring and OpenMP generically.'''
ctrans = Dynamo0p3ColourTrans()
otrans = DynamoOMPParallelLoopTrans()
const = LFRicConstants()
# Loop over all of the Invokes in the PSy object
for invoke in psy.invokes.invoke_list:
print("Transforming invoke '" + invoke.name + "'...")
schedule = invoke.schedule
# Colour all of the loops over cells unless they are on
# discontinuous spaces
for child in schedule.children:
if isinstance(child, Loop) \
and child.field_space.orig_name \
not in const.VALID_DISCONTINUOUS_NAMES \
and child.iteration_space == "cell_column":
ctrans.apply(child)
# Then apply OpenMP to each of the colour loops
for child in schedule.children:
if isinstance(child, Loop):
if child.loop_type == "colours":
otrans.apply(child.loop_body[0])
else:
otrans.apply(child)
schedule.view()
return psy
def test_loop_no_directive_dynamo0p3():
''' Test that applying Extract Transformation on a Loop without its
parent Directive when optimisations are applied in Dynamo0.3 API
raises a TransformationError. '''
etrans = DynamoExtractRegionTrans()
from psyclone.transformations import DynamoOMPParallelLoopTrans
# Test a Loop nested within the OMP Parallel DO Directive
_, invoke_info = parse(os.path.join(
DYNAMO_BASE_PATH, "4.13_multikernel_invokes_w3_anyd.f90"),
api=DYNAMO_API)
psy = PSyFactory(DYNAMO_API, distributed_memory=False).create(invoke_info)
invoke = psy.invokes.invoke_list[0]
schedule = invoke.schedule
# Apply DynamoOMPParallelLoopTrans to the second Loop
otrans = DynamoOMPParallelLoopTrans()
schedule, _ = otrans.apply(schedule.children[1])
loop = schedule.children[1].children[0]
# Try extracting the Loop inside the OMP Parallel DO region
with pytest.raises(TransformationError) as excinfo:
_, _ = etrans.apply(loop)
assert ("Extraction of a Loop without its parent Directive is not "
"allowed.") in str(excinfo)
def test_restrict_prolong_chain_anyd(tmpdir):
''' Test that we generate correct code for an invoke containing a
chain of discontinuous restrictions and continuous prolongations.
'''
_, invoke_info = parse(os.path.join(BASE_PATH,
"22.2.1_intergrid_3levels_anyd.f90"),
api=API)
psy = PSyFactory(API, distributed_memory=True).create(invoke_info)
schedule = psy.invokes.invoke_list[0].schedule
output = str(psy.gen)
# Check maps for any_discontinuous_space
expected = (
" map_adspc1_fld_m => fld_m_proxy%vspace%get_whole_dofmap()\n"
" map_adspc2_fld_f => fld_f_proxy%vspace%get_whole_dofmap()\n"
" map_adspc1_fld_c => fld_c_proxy%vspace%get_whole_dofmap()\n"
" map_adspc2_fld_m => fld_m_proxy%vspace%get_whole_dofmap()\n")
assert expected in output
# Check ndf and undf initialisations the second restrict kernel
# (fld_m to fld_c)
expected = (" ! Initialise number of DoFs for adspc1_fld_c\n"
" !\n"
" ndf_adspc1_fld_c = fld_c_proxy%vspace%get_ndf()\n"
" undf_adspc1_fld_c = fld_c_proxy%vspace%get_undf()\n"
" !\n"
" ! Initialise number of DoFs for adspc2_fld_m\n"
" !\n"
" ndf_adspc2_fld_m = fld_m_proxy%vspace%get_ndf()\n"
" undf_adspc2_fld_m = fld_m_proxy%vspace%get_undf()\n")
assert expected in output
# Check an example of restrict loop and all upper loop bounds
expected = (
" ! Call kernels and communication routines\n"
" !\n"
" DO cell=1,mesh_fld_m%get_last_edge_cell()\n"
" !\n"
" CALL restrict_kernel_code(nlayers, cell_map_fld_m(:,cell), "
"ncpc_fld_f_fld_m, ncell_fld_f, fld_m_proxy%data, fld_f_proxy%data, "
"undf_adspc1_fld_m, map_adspc1_fld_m(:,cell), ndf_adspc2_fld_f, "
"undf_adspc2_fld_f, map_adspc2_fld_f)\n"
" END DO\n")
assert expected in output
assert "DO cell=1,mesh_fld_c%get_last_edge_cell()" in output
assert "DO cell=1,mesh_fld_c%get_last_halo_cell(1)" in output
assert "DO cell=1,mesh_fld_m%get_last_halo_cell(1)" in output
# Check compilation
assert LFRicBuild(tmpdir).code_compiles(psy)
# Now do some transformations
from psyclone.transformations import Dynamo0p3ColourTrans, \
DynamoOMPParallelLoopTrans, TransformationError
otrans = DynamoOMPParallelLoopTrans()
ctrans = Dynamo0p3ColourTrans()
# Apply OMP to the first restrict kernel
_, _ = otrans.apply(schedule.children[0])
# Apply colouring and OMP to the first prolong kernel
_, _ = ctrans.apply(schedule.children[4])
_, _ = otrans.apply(schedule.children[4].loop_body[0])
output = str(psy.gen)
expected = (" !$omp parallel do default(shared), private(cell), "
"schedule(static)\n"
" DO cell=1,mesh_fld_m%get_last_edge_cell()\n"
" !\n"
" CALL restrict_kernel_code")
assert expected in output
expected = (" DO colour=1,ncolour_fld_m\n"
" !$omp parallel do default(shared), private(cell), "
"schedule(static)\n"
" DO cell=1,mesh_fld_c%get_last_halo_cell_per_colour"
"(colour,1)\n"
" !\n"
" CALL prolong_test_kernel_code")
assert expected in output
# Try to apply colouring to the second restrict kernel
with pytest.raises(TransformationError) as excinfo:
_, _ = ctrans.apply(schedule.children[1])
assert ("Loops iterating over a discontinuous function space "
"are not currently supported." in str(excinfo.value))
def test_extract_colouring_omp_dynamo0p3():
''' Test that extraction of a Kernel in an Invoke after applying
colouring and OpenMP optimisations produces the correct result
in Dynamo0.3 API. '''
const = LFRicConstants()
etrans = LFRicExtractTrans()
ctrans = Dynamo0p3ColourTrans()
otrans = DynamoOMPParallelLoopTrans()
psy, invoke = get_invoke("4.8_multikernel_invokes.f90",
DYNAMO_API,
idx=0,
dist_mem=False)
schedule = invoke.schedule
# First colour all of the loops over cells unless they are on
# discontinuous spaces
for child in schedule.children:
if isinstance(child, Loop) and child.field_space.orig_name \
not in const.VALID_DISCONTINUOUS_NAMES \
and child.iteration_space == "cell_column":
ctrans.apply(child)
# Then apply OpenMP to each of the colour loops
for child in schedule.children:
if isinstance(child, Loop):
if child.loop_type == "colours":
otrans.apply(child.loop_body[0])
else:
otrans.apply(child)
# Extract the second instance of ru_kernel_type after colouring
# and OpenMP are applied
child = schedule.children[2]
etrans.apply(child)
code = str(psy.gen)
output = ("""
! ExtractStart
!
CALL extract_psy_data%PreStart("multikernel_invokes_7_psy", """
""""invoke_0:ru_code:r0", 26, 3)
CALL extract_psy_data%PreDeclareVariable("a", a)
CALL extract_psy_data%PreDeclareVariable("b", b)
CALL extract_psy_data%PreDeclareVariable("basis_w0_qr", basis_w0_qr)
CALL extract_psy_data%PreDeclareVariable("basis_w2_qr", basis_w2_qr)
CALL extract_psy_data%PreDeclareVariable("basis_w3_qr", basis_w3_qr)
CALL extract_psy_data%PreDeclareVariable("c", c)
CALL extract_psy_data%PreDeclareVariable("cmap", cmap)
CALL extract_psy_data%PreDeclareVariable("diff_basis_w0_qr", """
"""diff_basis_w0_qr)
CALL extract_psy_data%PreDeclareVariable("diff_basis_w2_qr", """
"""diff_basis_w2_qr)
CALL extract_psy_data%PreDeclareVariable("e", e)
CALL extract_psy_data%PreDeclareVariable("istp", istp)
CALL extract_psy_data%PreDeclareVariable("map_w0", map_w0)
CALL extract_psy_data%PreDeclareVariable("map_w2", map_w2)
CALL extract_psy_data%PreDeclareVariable("map_w3", map_w3)
CALL extract_psy_data%PreDeclareVariable("ndf_w0", ndf_w0)
CALL extract_psy_data%PreDeclareVariable("ndf_w2", ndf_w2)
CALL extract_psy_data%PreDeclareVariable("ndf_w3", ndf_w3)
CALL extract_psy_data%PreDeclareVariable("nlayers", nlayers)
CALL extract_psy_data%PreDeclareVariable("np_xy_qr", np_xy_qr)
CALL extract_psy_data%PreDeclareVariable("np_z_qr", np_z_qr)
CALL extract_psy_data%PreDeclareVariable("rdt", rdt)
CALL extract_psy_data%PreDeclareVariable("undf_w0", undf_w0)
CALL extract_psy_data%PreDeclareVariable("undf_w2", undf_w2)
CALL extract_psy_data%PreDeclareVariable("undf_w3", undf_w3)
CALL extract_psy_data%PreDeclareVariable("weights_xy_qr", weights_xy_qr)
CALL extract_psy_data%PreDeclareVariable("weights_z_qr", weights_z_qr)
CALL extract_psy_data%PreDeclareVariable("b_post", b)
CALL extract_psy_data%PreDeclareVariable("cell_post", cell)
CALL extract_psy_data%PreDeclareVariable("colour_post", colour)
CALL extract_psy_data%PreEndDeclaration
CALL extract_psy_data%ProvideVariable("a", a)
CALL extract_psy_data%ProvideVariable("b", b)
CALL extract_psy_data%ProvideVariable("basis_w0_qr", basis_w0_qr)
CALL extract_psy_data%ProvideVariable("basis_w2_qr", basis_w2_qr)
CALL extract_psy_data%ProvideVariable("basis_w3_qr", basis_w3_qr)
CALL extract_psy_data%ProvideVariable("c", c)
CALL extract_psy_data%ProvideVariable("cmap", cmap)
CALL extract_psy_data%ProvideVariable("diff_basis_w0_qr", """
"""diff_basis_w0_qr)
CALL extract_psy_data%ProvideVariable("diff_basis_w2_qr", """
"""diff_basis_w2_qr)
CALL extract_psy_data%ProvideVariable("e", e)
CALL extract_psy_data%ProvideVariable("istp", istp)
CALL extract_psy_data%ProvideVariable("map_w0", map_w0)
CALL extract_psy_data%ProvideVariable("map_w2", map_w2)
CALL extract_psy_data%ProvideVariable("map_w3", map_w3)
CALL extract_psy_data%ProvideVariable("ndf_w0", ndf_w0)
CALL extract_psy_data%ProvideVariable("ndf_w2", ndf_w2)
CALL extract_psy_data%ProvideVariable("ndf_w3", ndf_w3)
CALL extract_psy_data%ProvideVariable("nlayers", nlayers)
CALL extract_psy_data%ProvideVariable("np_xy_qr", np_xy_qr)
CALL extract_psy_data%ProvideVariable("np_z_qr", np_z_qr)
CALL extract_psy_data%ProvideVariable("rdt", rdt)
CALL extract_psy_data%ProvideVariable("undf_w0", undf_w0)
CALL extract_psy_data%ProvideVariable("undf_w2", undf_w2)
CALL extract_psy_data%ProvideVariable("undf_w3", undf_w3)
CALL extract_psy_data%ProvideVariable("weights_xy_qr", weights_xy_qr)
CALL extract_psy_data%ProvideVariable("weights_z_qr", weights_z_qr)
CALL extract_psy_data%PreEnd
DO colour=1,ncolour
!$omp parallel do default(shared), private(cell), schedule(static)
DO cell=1,mesh%get_last_edge_cell_per_colour(colour)
!
CALL ru_code(nlayers, b_proxy%data, a_proxy%data, istp, rdt, """
"c_proxy%data, e_proxy(1)%data, e_proxy(2)%data, "
"e_proxy(3)%data, ndf_w2, undf_w2, "
"map_w2(:,cmap(colour, cell)), "
"basis_w2_qr, diff_basis_w2_qr, ndf_w3, undf_w3, "
"map_w3(:,cmap(colour, cell)), basis_w3_qr, ndf_w0, undf_w0, "
"map_w0(:,cmap(colour, cell)), basis_w0_qr, diff_basis_w0_qr, "
"""np_xy_qr, np_z_qr, weights_xy_qr, weights_z_qr)
END DO
!$omp end parallel do
END DO
CALL extract_psy_data%PostStart
CALL extract_psy_data%ProvideVariable("b_post", b)
CALL extract_psy_data%ProvideVariable("cell_post", cell)
CALL extract_psy_data%ProvideVariable("colour_post", colour)
CALL extract_psy_data%PostEnd
!
! ExtractEnd""")
assert output in code
def test_extract_single_builtin_dynamo0p3():
''' Test that extraction of a BuiltIn in an Invoke produces the
correct result in Dynamo0.3 API without and with optimisations.
'''
etrans = LFRicExtractTrans()
otrans = DynamoOMPParallelLoopTrans()
psy, invoke = get_invoke("15.1.2_builtin_and_normal_kernel_invoke.f90",
DYNAMO_API,
idx=0,
dist_mem=False)
schedule = invoke.schedule
etrans.apply(schedule.children[1])
code = str(psy.gen)
output = """! ExtractStart
!
CALL extract_psy_data%PreStart("single_invoke_builtin_then_kernel_psy", """ \
""""invoke_0:setval_c:r0", 0, 1)
CALL extract_psy_data%PreDeclareVariable("df_post", df)
CALL extract_psy_data%PreEndDeclaration
CALL extract_psy_data%PreEnd
DO df=1,undf_aspc1_f2
f2_proxy%data(df) = 0.0
END DO
CALL extract_psy_data%PostStart
CALL extract_psy_data%ProvideVariable("df_post", df)
CALL extract_psy_data%PostEnd
!
! ExtractEnd"""
assert output in code
# Test extract with OMP Parallel optimisation
psy, invoke = get_invoke("15.1.1_builtin_and_normal_kernel_invoke_2.f90",
DYNAMO_API,
idx=0,
dist_mem=False)
schedule = invoke.schedule
otrans.apply(schedule.children[1])
etrans.apply(schedule.children[1])
code_omp = str(psy.gen)
output = """
! ExtractStart
!
CALL extract_psy_data%PreStart("single_invoke_psy", """ \
""""invoke_0:inc_ax_plus_y:r0", 0, 1)
CALL extract_psy_data%PreDeclareVariable("df_post", df)
CALL extract_psy_data%PreEndDeclaration
CALL extract_psy_data%PreEnd
!$omp parallel do default(shared), private(df), schedule(static)
DO df=1,undf_aspc1_f1
f1_proxy%data(df) = 0.5_r_def*f1_proxy%data(df) + f2_proxy%data(df)
END DO
!$omp end parallel do
CALL extract_psy_data%PostStart
CALL extract_psy_data%ProvideVariable("df_post", df)
CALL extract_psy_data%PostEnd
!
! ExtractEnd"""
assert output in code_omp
# TODO #637 (no builtin support)
# Builtins are not yet support, so some arguments are missing. This
# is an excerpt of missing lines, which will cause this test to x-fail
not_yet_working = [
'CALL extract_psy_data%ProvideVariable("f1", f1)',
'CALL extract_psy_data%ProvideVariable("f2", f2)'
]
for line in not_yet_working:
if line not in code:
pytest.xfail("#637 LFRic builtins are not supported yet.")
if line not in code_omp:
pytest.xfail("#637 LFRic builtins are not supported yet.")
assert False, "X-failing test suddenly working: #637 lfric builtins."
def test_extract_colouring_omp_dynamo0p3(tmpdir):
''' Test that extraction of a Kernel in an Invoke after applying
colouring and OpenMP optimisations produces the correct result
in Dynamo0.3 API. '''
from psyclone.transformations import Dynamo0p3ColourTrans, \
DynamoOMPParallelLoopTrans
from psyclone.dynamo0p3 import VALID_DISCONTINUOUS_FUNCTION_SPACE_NAMES
etrans = DynamoExtractRegionTrans()
ctrans = Dynamo0p3ColourTrans()
otrans = DynamoOMPParallelLoopTrans()
_, invoke_info = parse(os.path.join(DYNAMO_BASE_PATH,
"4.8_multikernel_invokes.f90"),
api=DYNAMO_API)
psy = PSyFactory(DYNAMO_API, distributed_memory=False).create(invoke_info)
invoke = psy.invokes.invoke_list[0]
schedule = invoke.schedule
# First colour all of the loops over cells unless they are on
# discontinuous spaces
cschedule = schedule
for child in schedule.children:
if isinstance(child, Loop) and child.field_space.orig_name \
not in VALID_DISCONTINUOUS_FUNCTION_SPACE_NAMES \
and child.iteration_space == "cells":
cschedule, _ = ctrans.apply(child)
# Then apply OpenMP to each of the colour loops
schedule = cschedule
for child in schedule.children:
if isinstance(child, Loop):
if child.loop_type == "colours":
schedule, _ = otrans.apply(child.loop_body[0])
else:
schedule, _ = otrans.apply(child)
# Extract the second instance of ru_kernel_type after colouring
# and OpenMP are applied
child = schedule.children[2]
schedule, _ = etrans.apply(child)
code = str(psy.gen)
output = (
" ! ExtractStart\n"
" ! CALL write_extract_arguments(argument_list)\n"
" !\n"
" DO colour=1,ncolour\n"
" !$omp parallel do default(shared), private(cell), "
"schedule(static)\n"
" DO cell=1,mesh%get_last_edge_cell_per_colour(colour)\n"
" !\n"
" CALL ru_code(nlayers, b_proxy%data, a_proxy%data, istp, "
"rdt, c_proxy%data, e_proxy(1)%data, e_proxy(2)%data, "
"e_proxy(3)%data, ndf_w2, undf_w2, map_w2(:,cmap(colour, cell)), "
"basis_w2_qr, diff_basis_w2_qr, ndf_w3, undf_w3, "
"map_w3(:,cmap(colour, cell)), basis_w3_qr, ndf_w0, undf_w0, "
"map_w0(:,cmap(colour, cell)), basis_w0_qr, diff_basis_w0_qr, "
"np_xy_qr, np_z_qr, weights_xy_qr, weights_z_qr)\n"
" END DO \n"
" !$omp end parallel do\n"
" END DO \n"
" !\n"
" ! ExtractEnd\n")
assert output in code
assert Dynamo0p3Build(tmpdir).code_compiles(psy)