def test_1kern_trans(kernel_outputdir):
''' Check that we generate the correct code when an invoke contains
the same kernel more than once but only one of them is transformed. '''
psy, invoke = get_invoke("4_multikernel_invokes.f90",
api="dynamo0.3",
idx=0)
sched = invoke.schedule
kernels = sched.coded_kernels()
# We will transform the second kernel but not the first
kern = kernels[1]
rtrans = ACCRoutineTrans()
_, _ = rtrans.apply(kern)
# Generate the code (this triggers the generation of a new kernel)
code = str(psy.gen).lower()
tag = re.search('use testkern(.+?)_mod', code).group(1)
# We should have a USE for the original kernel and a USE for the new one
assert "use testkern{0}_mod, only: testkern{0}_code".format(tag) in code
assert "use testkern_mod, only: testkern_code" in code
# Similarly, we should have calls to both the original and new kernels
assert "call testkern_code(" in code
assert "call testkern{0}_code(".format(tag) in code
first = code.find("call testkern_code(")
second = code.find("call testkern{0}_code(".format(tag))
assert first < second
assert Dynamo0p3Build(kernel_outputdir).code_compiles(psy)
def test_2kern_trans(kernel_outputdir):
''' Check that we generate correct code when we transform two kernels
within a single invoke. '''
psy, invoke = get_invoke("4.5.2_multikernel_invokes.f90",
api="dynamo0.3",
idx=0)
sched = invoke.schedule
kernels = sched.walk(Kern)
assert len(kernels) == 5
ktrans = Dynamo0p3KernelConstTrans()
_, _ = ktrans.apply(kernels[1], number_of_layers=100)
_, _ = ktrans.apply(kernels[2], number_of_layers=100)
# Generate the code (this triggers the generation of new kernels)
code = str(psy.gen).lower()
# Find the tags added to the kernel/module names
for match in re.finditer('use testkern_any_space_2(.+?)_mod', code):
tag = match.group(1)
assert ("use testkern_any_space_2{0}_mod, only: "
"testkern_any_space_2{0}_code".format(tag) in code)
assert "call testkern_any_space_2{0}_code(".format(tag) in code
filepath = os.path.join(str(kernel_outputdir),
"testkern_any_space_2{0}_mod.f90".format(tag))
assert os.path.isfile(filepath)
assert "nlayers = 100" in open(filepath).read()
assert "use testkern_any_space_2_mod, only" not in code
assert "call testkern_any_space_2_code(" not in code
assert Dynamo0p3Build(kernel_outputdir).code_compiles(psy)
def test_operator_nofield_scalar_deref(tmpdir, dist_mem):
''' Tests that an operator with no field and a
scalar argument is implemented correctly in the PSy layer when both
are obtained by dereferencing derived type objects '''
_, invoke_info = parse(os.path.join(
BASE_PATH, "10.6.1_operator_no_field_scalar_deref.f90"),
api=TEST_API)
psy = PSyFactory(TEST_API, distributed_memory=dist_mem).create(invoke_info)
gen = str(psy.gen)
assert Dynamo0p3Build(tmpdir).code_compiles(psy)
if dist_mem:
assert "mesh => opbox_my_mapping_proxy%fs_from%get_mesh()" in gen
assert "nlayers = opbox_my_mapping_proxy%fs_from%get_nlayers()" in gen
assert "ndf_w2 = opbox_my_mapping_proxy%fs_from%get_ndf()" in gen
assert ("qr_get_instance%compute_function(BASIS, "
"opbox_my_mapping_proxy%fs_from, "
"dim_w2, ndf_w2, basis_w2_qr_get_instance)" in gen)
if dist_mem:
assert "DO cell=1,mesh%get_last_halo_cell(1)" in gen
else:
assert ("DO cell=1,opbox_my_mapping_proxy%fs_from%get_ncell()" in gen)
assert ("(cell, nlayers, opbox_my_mapping_proxy%ncell_3d, "
"opbox_my_mapping_proxy%local_stencil, box_b, ndf_w2, "
"basis_w2_qr_get_instance, np_xy_qr_get_instance, "
"np_z_qr_get_instance, weights_xy_qr_get_instance,"
" weights_z_qr_get_instance)" in gen)
def test_op_orient_different_space(tmpdir):
''' Tests that an operator on different spaces requiring orientation
information is implemented correctly in the PSy layer '''
_, invoke_info = parse(os.path.join(
BASE_PATH, "10.4_operator_orient_different_"
"space.f90"),
api=TEST_API)
psy = PSyFactory(TEST_API, distributed_memory=True).create(invoke_info)
gen_str = str(psy.gen)
assert Dynamo0p3Build(tmpdir).code_compiles(psy)
assert (
"INTEGER, pointer :: orientation_w1(:) => null(), orientation_w2(:)"
" => null()" in gen_str)
assert "ndf_w2 = my_mapping_proxy%fs_from%get_ndf()" in gen_str
assert "ndf_w1 = my_mapping_proxy%fs_to%get_ndf()" in gen_str
assert "dim_w1 = my_mapping_proxy%fs_to%get_dim_space()" in gen_str
assert ("CALL qr%compute_function(BASIS, my_mapping_proxy%fs_to, "
"dim_w1, ndf_w1, basis_w1_qr)" in gen_str)
assert ("orientation_w2 => my_mapping_proxy%fs_from%get_cell_orientation("
"cell)" in gen_str)
assert (
"orientation_w1 => my_mapping_proxy%fs_to%get_cell_orientation(cell)"
in gen_str)
assert ("(cell, nlayers, my_mapping_proxy%ncell_3d, "
"my_mapping_proxy%local_stencil, chi_proxy(1)%data, "
"chi_proxy(2)%data, chi_proxy(3)%data, ndf_w1, basis_w1_qr, "
"orientation_w1, ndf_w2, orientation_w2, ndf_w0, undf_w0, "
"map_w0(:,cell), diff_basis_w0_qr, np_xy_qr, np_z_qr, "
"weights_xy_qr, weights_z_qr)" in gen_str)
def test_prolong_vector(tmpdir):
''' Check that we generate correct code when an inter-grid kernel
takes a field vector as argument '''
_, invoke_info = parse(os.path.join(BASE_PATH,
"22.4_intergrid_prolong_vec.f90"),
api=API)
psy = PSyFactory(API, distributed_memory=True).create(invoke_info)
output = str(psy.gen)
assert Dynamo0p3Build(tmpdir).code_compiles(psy)
assert "TYPE(field_type), intent(inout) :: field1(3)" in output
assert "TYPE(field_proxy_type) field1_proxy(3)" in output
# Make sure we always index into the field arrays
assert " field1%" not in output
assert " field2%" not in output
assert ("ncpc_field1_field2, ncell_field1, field1_proxy(1)%data, "
"field1_proxy(2)%data, field1_proxy(3)%data, field2_proxy(1)%data,"
" field2_proxy(2)%data, field2_proxy(3)%data, ndf_w1" in output)
for idx in [1, 2, 3]:
assert (" IF (field2_proxy({0})%is_dirty(depth=1)) THEN\n"
" CALL field2_proxy({0})%halo_exchange(depth=1)\n"
" END IF \n".format(idx) in output)
assert ("field1_proxy({0}) = field1({0})%get_proxy()".format(idx)
in output)
assert "CALL field1_proxy({0})%set_dirty()".format(idx) in output
assert "CALL field1_proxy({0})%set_clean(1)".format(idx) in output
def test_operator_nofield(tmpdir):
''' Tests that an operator with no field on the same space is
implemented correctly in the PSy layer '''
_, invoke_info = parse(os.path.join(BASE_PATH,
"10.1_operator_nofield.f90"),
api=TEST_API)
psy = PSyFactory(TEST_API, distributed_memory=True).create(invoke_info)
gen_code_str = str(psy.gen)
assert Dynamo0p3Build(tmpdir).code_compiles(psy)
assert (
"SUBROUTINE invoke_0_testkern_operator_nofield_type(mm_w2, chi, qr)"
in gen_code_str)
assert "TYPE(operator_type), intent(inout) :: mm_w2" in gen_code_str
assert "TYPE(operator_proxy_type) mm_w2_proxy" in gen_code_str
assert "mm_w2_proxy = mm_w2%get_proxy()" in gen_code_str
assert "undf_w2" not in gen_code_str
assert "map_w2" not in gen_code_str
assert ("CALL testkern_operator_nofield_code(cell, nlayers, "
"mm_w2_proxy%ncell_3d, mm_w2_proxy%local_stencil, "
"chi_proxy(1)%data, chi_proxy(2)%data, chi_proxy(3)%data, "
"ndf_w2, basis_w2_qr, ndf_w0, undf_w0, "
"map_w0(:,cell), diff_basis_w0_qr, np_xy_qr, np_z_qr, "
"weights_xy_qr, weights_z_qr)" in gen_code_str)
def test_extract_kernel_and_builtin_dynamo0p3(tmpdir):
''' Test that extraction of a Kernel and a BuiltIny in an Invoke
produces the correct result in Dynamo0.3 API. '''
etrans = DynamoExtractRegionTrans()
_, 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:3])
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"
" DO cell=1,f3_proxy%vspace%get_ncell()\n"
" !\n"
" CALL testkern_code_w2_only(nlayers, f3_proxy%data, "
"f2_proxy%data, ndf_w2, undf_w2, map_w2(:,cell))\n"
" END DO \n"
" !\n"
" ! ExtractEnd\n")
assert output in code
assert Dynamo0p3Build(tmpdir).code_compiles(psy)
def test_gh_inc_nohex_1(tmpdir, monkeypatch):
'''If COMPUTE_ANNEXED_DOFS is True, then a gh_inc access to a field in
a kernel (iterating to the l1 halo) does not require a halo
exchange when the previous writer is known and iterates over dofs
to nannexed, halo(1), or halo max depth
'''
# ensure that COMPUTE_ANNEXED_DOFS is True
config = Config.get()
dyn_config = config.api_conf(API)
monkeypatch.setattr(dyn_config, "_compute_annexed_dofs", True)
# parse and get psy schedule
_, info = parse(os.path.join(BASE_PATH,
"14.12_halo_wdofs_to_inc.f90"),
api=API)
psy = PSyFactory(API, distributed_memory=True).create(info)
schedule = psy.invokes.invoke_list[0].schedule
def check_schedule(schedule):
'''Check this schedule has expected structure (loop, haloexchange,
loop). In paricular there should be no halo exchange for the
write-to-gh_inc dependence.
:param schedule: a dynamo0.3 API schedule object
:type schedule: :py:class:`psyclone.dynamo0p3.DynInvokeSchedule`.
'''
assert len(schedule.children) == 3
loop1 = schedule.children[0]
haloex = schedule.children[1]
loop2 = schedule.children[2]
assert isinstance(loop1, DynLoop)
assert isinstance(haloex, DynHaloExchange)
assert haloex.field.name == "f2"
assert haloex.required() == (True, False)
assert isinstance(loop2, DynLoop)
# 1st loop should iterate over dofs to nannexed. Check output
assert schedule.children[0].upper_bound_name == "nannexed"
check_schedule(schedule)
# just check compilation here (not later in this test) as
# compilation of redundant computation is checked separately
assert Dynamo0p3Build(tmpdir).code_compiles(psy)
# make 1st loop iterate over dofs to the level 1 halo and check output
rc_trans = Dynamo0p3RedundantComputationTrans()
rc_trans.apply(schedule.children[0], depth=1)
assert schedule.children[0].upper_bound_name == "dof_halo"
assert schedule.children[0].upper_bound_halo_depth == 1
check_schedule(schedule)
# make 1st loop iterate over dofs to the maximum halo depth and
# check output
rc_trans.apply(schedule.children[0])
assert schedule.children[0].upper_bound_name == "dof_halo"
assert not schedule.children[0].upper_bound_halo_depth
check_schedule(schedule)
def test_operator_nofield_scalar(tmpdir):
''' Tests that an operator with no field and a
scalar argument is implemented correctly in the PSy layer '''
_, invoke_info = parse(os.path.join(BASE_PATH,
"10.6_operator_no_field_scalar.f90"),
api=TEST_API)
psy = PSyFactory(TEST_API, distributed_memory=True).create(invoke_info)
gen = str(psy.gen)
assert Dynamo0p3Build(tmpdir).code_compiles(psy)
assert "mesh => my_mapping_proxy%fs_from%get_mesh()" in gen
assert "nlayers = my_mapping_proxy%fs_from%get_nlayers()" in gen
assert "ndf_w2 = my_mapping_proxy%fs_from%get_ndf()" in gen
assert "DO cell=1,mesh%get_last_halo_cell(1)" in gen
assert ("(cell, nlayers, my_mapping_proxy%ncell_3d, my_mapping_proxy%"
"local_stencil, b, ndf_w2, basis_w2_qr, np_xy_qr, np_z_qr, "
"weights_xy_qr, weights_z_qr)" in gen)
def test_field_qr_deref(tmpdir):
''' Tests that a call, with a set of fields requiring
quadrature, produces correct code when the quadrature is supplied as the
component of a derived type. '''
_, invoke_info = parse(os.path.join(BASE_PATH,
"1.1.1_single_invoke_qr_deref.f90"),
api="dynamo0.3")
for dist_mem in [True, False]:
psy = PSyFactory("dynamo0.3",
distributed_memory=dist_mem).create(invoke_info)
assert Dynamo0p3Build(tmpdir).code_compiles(psy)
gen = str(psy.gen)
print(gen)
assert (
" SUBROUTINE invoke_0_testkern_qr_type(f1, f2, m1, a, m2, istp,"
" qr_data)\n" in gen)
assert "TYPE(quadrature_xyoz_type), intent(in) :: qr_data" in gen
def test_node_list_error(tmpdir):
''' Test that applying Extract Transformation on objects which are not
Nodes or a list of Nodes raises a TransformationError. Also raise
transformation errors when the Nodes do not have the same parent
if they are incorrectly ordered. '''
etrans = DynamoExtractRegionTrans()
_, invoke_info = parse(os.path.join(
DYNAMO_BASE_PATH, "3.2_multi_functions_multi_named_invokes.f90"),
api=DYNAMO_API)
psy = PSyFactory(DYNAMO_API, distributed_memory=False).create(invoke_info)
invoke0 = psy.invokes.invoke_list[0]
invoke1 = psy.invokes.invoke_list[1]
# Supply an object which is not a Node or a list of Nodes
with pytest.raises(TransformationError) as excinfo:
etrans.apply(invoke0)
assert ("Error in DynamoExtractRegionTrans: Argument must be "
"a single Node in a Schedule or a list of Nodes in a Schedule "
"but have been passed an object of type: "
"<class 'psyclone.dynamo0p3.DynInvoke'>") in str(excinfo)
# Supply Nodes in incorrect order or duplicate Nodes
node_list = [
invoke0.schedule.children[0], invoke0.schedule.children[0],
invoke0.schedule.children[1]
]
with pytest.raises(TransformationError) as excinfo:
etrans.apply(node_list)
assert "Children are not consecutive children of one parent:" \
in str(excinfo)
assert "has position 0, but previous child had position 0."\
in str(excinfo)
# Supply Nodes which are not children of the same parent
node_list = [
invoke0.schedule.children[1], invoke1.schedule.children[0],
invoke0.schedule.children[2]
]
with pytest.raises(TransformationError) as excinfo:
etrans.apply(node_list)
assert ("supplied nodes are not children of the same "
"parent.") in str(excinfo)
assert Dynamo0p3Build(tmpdir).code_compiles(psy)
def test_operator_nofield_different_space(tmpdir):
''' tests that an operator with no field on different spaces is
implemented correctly in the PSy layer '''
_, invoke_info = parse(os.path.join(
BASE_PATH, "10.5_operator_no_field_different_"
"space.f90"),
api=TEST_API)
psy = PSyFactory(TEST_API, distributed_memory=True).create(invoke_info)
gen = str(psy.gen)
assert Dynamo0p3Build(tmpdir).code_compiles(psy)
assert "mesh => my_mapping_proxy%fs_from%get_mesh()" in gen
assert "nlayers = my_mapping_proxy%fs_from%get_nlayers()" in gen
assert "ndf_w3 = my_mapping_proxy%fs_from%get_ndf()" in gen
assert "ndf_w2 = my_mapping_proxy%fs_to%get_ndf()" in gen
# We compute operators redundantly (out to the L1 halo)
assert "DO cell=1,mesh%get_last_halo_cell(1)" in gen
assert ("(cell, nlayers, my_mapping_proxy%ncell_3d, my_mapping_proxy%"
"local_stencil, ndf_w2, ndf_w3)" in gen)
def test_operator_bc_kernel(tmpdir):
''' Tests that a kernel with a particular name is recognised as a
kernel that applies boundary conditions to operators and that
appropriate code is added to support this. '''
_, invoke_info = parse(os.path.join(BASE_PATH,
"12.4_enforce_op_bc_kernel.f90"),
api=TEST_API)
psy = PSyFactory(TEST_API, distributed_memory=True).create(invoke_info)
generated_code = str(psy.gen)
output1 = "INTEGER, pointer :: boundary_dofs_op_a(:,:) => null()"
assert output1 in generated_code
output2 = "boundary_dofs_op_a => op_a_proxy%fs_to%get_boundary_dofs()"
assert output2 in generated_code
output3 = (
"CALL enforce_operator_bc_code(cell, nlayers, op_a_proxy%ncell_3d, "
"op_a_proxy%local_stencil, ndf_any_space_1_op_a, "
"ndf_any_space_2_op_a, boundary_dofs_op_a)")
assert output3 in generated_code
assert Dynamo0p3Build(tmpdir).code_compiles(psy)
def test_operator_deref(tmpdir, dist_mem):
''' Tests that we generate correct names for an operator in the PSy
layer when obtained by de-referencing a derived type in the Algorithm
layer '''
_, invoke_info = parse(os.path.join(BASE_PATH, "10.8_operator_deref.f90"),
api=TEST_API)
psy = PSyFactory(TEST_API, distributed_memory=dist_mem).create(invoke_info)
generated_code = str(psy.gen)
assert Dynamo0p3Build(tmpdir).code_compiles(psy)
assert ("SUBROUTINE invoke_0_testkern_operator_type(mm_w0_op, chi, a, qr)"
in generated_code)
assert "TYPE(operator_type), intent(inout) :: mm_w0_op" in generated_code
assert "TYPE(operator_proxy_type) mm_w0_op_proxy" in generated_code
assert "mm_w0_op_proxy = mm_w0_op%get_proxy()" in generated_code
assert ("CALL testkern_operator_code(cell, nlayers, "
"mm_w0_op_proxy%ncell_3d, mm_w0_op_proxy%local_stencil, "
"chi_proxy(1)%data, chi_proxy(2)%data, chi_proxy(3)%data, a, "
"ndf_w0, undf_w0, map_w0(:,cell), basis_w0_qr, "
"diff_basis_w0_qr, np_xy_qr, np_z_qr, weights_xy_qr, "
"weights_z_qr)" in generated_code)
def test_operator_orientation(tmpdir):
''' Tests that an operator requiring orientation information is
implemented correctly in the PSy layer '''
_, invoke_info = parse(os.path.join(BASE_PATH, "10.2_operator_orient.f90"),
api=TEST_API)
psy = PSyFactory(TEST_API, distributed_memory=True).create(invoke_info)
gen_str = str(psy.gen)
assert Dynamo0p3Build(tmpdir).code_compiles(psy)
assert ("SUBROUTINE invoke_0_testkern_operator_orient_type(mm_w1, chi, qr)"
in gen_str)
assert "TYPE(operator_type), intent(inout) :: mm_w1" in gen_str
assert "TYPE(operator_proxy_type) mm_w1_proxy" in gen_str
assert "mm_w1_proxy = mm_w1%get_proxy()" in gen_str
assert ("orientation_w1 => mm_w1_proxy%fs_from%get_cell_orientation"
"(cell)" in gen_str)
assert ("CALL testkern_operator_orient_code(cell, nlayers, "
"mm_w1_proxy%ncell_3d, mm_w1_proxy%local_stencil, "
"chi_proxy(1)%data, chi_proxy(2)%data, chi_proxy(3)%data, ndf_w1, "
"basis_w1_qr, orientation_w1, ndf_w0, undf_w0, map_w0(:,cell), "
"diff_basis_w0_qr, np_xy_qr, np_z_qr, weights_xy_qr, "
"weights_z_qr)" in gen_str)
def infra_compile(tmpdir_factory, request):
'''A per-session initialisation function that sets the compilation flags
in the Compile class based on command line options for --compile,
--compileopencl, --f90, --f90flags. Then makes sure that the
infrastructure files for the dynamo0p3 and gocean1p0 APIs are compiled
(if compilation was enabled).
'''
from psyclone.tests.utilities import Compile
Compile.store_compilation_flags(request.config)
from psyclone.tests.dynamo0p3_build import Dynamo0p3Build
# Create a temporary directory to store the compiled files.
# Note that this directory is unique even if compiled in
# parallel, i.e. each process has its own copy of the
# compiled infrastructure file, which avoids the problem
# of synchronisation between the processes.
tmpdir = tmpdir_factory.mktemp('dynamo_wrapper')
# This is the first instance created. This will trigger
# compilation of the infrastructure files.
Dynamo0p3Build(tmpdir)
from psyclone.tests.gocean1p0_build import GOcean1p0Build
tmpdir = tmpdir_factory.mktemp('dl_esm_inf')
GOcean1p0Build(tmpdir)
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_any_discontinuous_space_1_fld_m => "
"fld_m_proxy%vspace%get_whole_dofmap()\n"
" map_any_discontinuous_space_2_fld_f => "
"fld_f_proxy%vspace%get_whole_dofmap()\n"
" map_any_discontinuous_space_1_fld_c => "
"fld_c_proxy%vspace%get_whole_dofmap()\n"
" map_any_discontinuous_space_2_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 "
"any_discontinuous_space_1_fld_c\n"
" !\n"
" ndf_any_discontinuous_space_1_fld_c = "
"fld_c_proxy%vspace%get_ndf()\n"
" undf_any_discontinuous_space_1_fld_c = "
"fld_c_proxy%vspace%get_undf()\n"
" !\n"
" ! Initialise number of DoFs for "
"any_discontinuous_space_2_fld_m\n"
" !\n"
" ndf_any_discontinuous_space_2_fld_m = "
"fld_m_proxy%vspace%get_ndf()\n"
" undf_any_discontinuous_space_2_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_any_discontinuous_space_1_fld_m, "
"map_any_discontinuous_space_1_fld_m(:,cell), "
"ndf_any_discontinuous_space_2_fld_f, "
"undf_any_discontinuous_space_2_fld_f, "
"map_any_discontinuous_space_2_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 Dynamo0p3Build(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_gh_inc_nohex_2(tmpdir, monkeypatch):
'''If COMPUTE_ANNEXED_DOFS is False, then a gh_inc access to a field in
a kernel (iterating to the l1 halo) does require a halo
exchange when the previous writer is known and iterates over dofs
to ndofs but does not if it iterates to halo(1), or halo max depth
'''
# ensure that COMPUTE_ANNEXED_DOFS is False
config = Config.get()
dyn_config = config.api_conf(API)
monkeypatch.setattr(dyn_config, "_compute_annexed_dofs", False)
# parse and get psy schedule
_, info = parse(os.path.join(BASE_PATH,
"14.12_halo_wdofs_to_inc.f90"),
api=API)
psy = PSyFactory(API, distributed_memory=True).create(info)
schedule = psy.invokes.invoke_list[0].schedule
# 1st loop should iterate over dofs to ndofs. Check output
loop1 = schedule.children[0]
haloex1 = schedule.children[1]
haloex2 = schedule.children[2]
loop2 = schedule.children[3]
assert len(schedule.children) == 4
assert isinstance(loop1, DynLoop)
assert loop1.upper_bound_name == "ndofs"
assert isinstance(haloex1, DynHaloExchange)
assert haloex1.field.name == "f1"
assert haloex1.required() == (True, True)
assert isinstance(haloex2, DynHaloExchange)
assert haloex2.field.name == "f2"
assert haloex2.required() == (True, False)
assert isinstance(loop2, DynLoop)
# just check compilation here (not later in this test) as
# compilation of redundant computation is checked separately
assert Dynamo0p3Build(tmpdir).code_compiles(psy)
# make 1st loop iterate over dofs to the level 1 halo and check
# output. There should be no halo exchange for field "f1"
rc_trans = Dynamo0p3RedundantComputationTrans()
rc_trans.apply(schedule.children[0], depth=1)
loop1 = schedule.children[0]
haloex = schedule.children[1]
loop2 = schedule.children[2]
assert len(schedule.children) == 3
assert isinstance(loop1, DynLoop)
assert loop1.upper_bound_name == "dof_halo"
assert loop1.upper_bound_halo_depth == 1
assert isinstance(haloex, DynHaloExchange)
assert haloex.field.name == "f2"
assert haloex.required() == (True, False)
assert isinstance(loop2, DynLoop)
# make 1st loop iterate over dofs to the maximum halo depth and
# check output
rc_trans.apply(schedule.children[0])
loop1 = schedule.children[0]
haloex = schedule.children[1]
loop2 = schedule.children[2]
assert len(schedule.children) == 3
assert isinstance(loop1, DynLoop)
assert loop1.upper_bound_name == "dof_halo"
assert not loop1.upper_bound_halo_depth
assert isinstance(haloex, DynHaloExchange)
assert haloex.field.name == "f2"
assert haloex.required() == (True, False)
assert isinstance(loop2, DynLoop)
def test_restrict_prolong_chain(tmpdir, dist_mem):
''' Test when we have a single invoke containing a chain of
restrictions and prolongations '''
_, invoke_info = parse(os.path.join(BASE_PATH,
"22.2_intergrid_3levels.f90"),
api=API)
psy = PSyFactory(API, distributed_memory=dist_mem).create(invoke_info)
output = str(psy.gen)
assert Dynamo0p3Build(tmpdir).code_compiles(psy)
expected = (
" ! Look-up mesh objects and loop limits for inter-grid "
"kernels\n"
" !\n"
" mesh_fld_f => fld_f_proxy%vspace%get_mesh()\n"
" mesh_fld_m => fld_m_proxy%vspace%get_mesh()\n"
" mmap_fld_f_fld_m => mesh_fld_m%get_mesh_map(mesh_fld_f)\n"
" cell_map_fld_m => mmap_fld_f_fld_m%get_whole_cell_map()\n")
assert expected in output
if dist_mem:
expected = (
" ncell_fld_f = mesh_fld_f%get_last_halo_cell(depth=2)\n"
" ncpc_fld_f_fld_m = mmap_fld_f_fld_m%"
"get_ntarget_cells_per_source_cell()\n"
" mesh_fld_c => fld_c_proxy%vspace%get_mesh()\n"
" mmap_fld_m_fld_c => mesh_fld_c%get_mesh_map(mesh_fld_m)\n"
" cell_map_fld_c => mmap_fld_m_fld_c%get_whole_cell_map()\n"
" ncell_fld_m = mesh_fld_m%get_last_halo_cell(depth=2)\n"
" ncpc_fld_m_fld_c = mmap_fld_m_fld_c%"
"get_ntarget_cells_per_source_cell()\n")
else:
expected = (
" ncell_fld_f = fld_f_proxy%vspace%get_ncell()\n"
" ncpc_fld_f_fld_m = mmap_fld_f_fld_m%"
"get_ntarget_cells_per_source_cell()\n"
" mesh_fld_c => fld_c_proxy%vspace%get_mesh()\n"
" mmap_fld_m_fld_c => mesh_fld_c%get_mesh_map(mesh_fld_m)\n"
" cell_map_fld_c => mmap_fld_m_fld_c%get_whole_cell_map()\n"
" ncell_fld_m = fld_m_proxy%vspace%get_ncell()\n"
" ncpc_fld_m_fld_c = mmap_fld_m_fld_c%get_ntarget_cells_"
"per_source_cell()\n")
assert expected in output
# Check that we haven't got duplicated output
assert output.count("mesh_fld_m => fld_m_proxy%vspace%get_mesh") == 1
assert output.count("ncell_fld_m = ") == 1
assert output.count("ncell_fld_f = ") == 1
if dist_mem:
# Have two potential halo exchanges before 1st prolong because
# of continuous "read"er and "inc" writer
expected = (" IF (fld_m_proxy%is_dirty(depth=1)) THEN\n"
" CALL fld_m_proxy%halo_exchange(depth=1)\n"
" END IF \n"
" !\n"
" IF (fld_c_proxy%is_dirty(depth=1)) THEN\n"
" CALL fld_c_proxy%halo_exchange(depth=1)\n"
" END IF \n"
" !\n"
" DO cell=1,mesh_fld_c%get_last_halo_cell(1)\n")
assert expected in output
# Since we loop into L1 halo of the coarse mesh, the L1 halo
# of the fine(r) mesh will now be clean. Therefore, no halo
# swap before the next prolongation required for fld_m
expected = (" ! Set halos dirty/clean for fields modified in the "
"above loop\n"
" !\n"
" CALL fld_m_proxy%set_dirty()\n"
" CALL fld_m_proxy%set_clean(1)\n"
" !\n"
" IF (fld_f_proxy%is_dirty(depth=1)) THEN\n"
" CALL fld_f_proxy%halo_exchange(depth=1)\n"
" END IF \n"
" !\n"
" DO cell=1,mesh_fld_m%get_last_halo_cell(1)\n")
assert expected in output
# Again the L1 halo for fld_f will now be clean but for restriction
# we need the L2 halo to be clean. There's a set_clean(1) for
# fld_f because the above loop over the coarser fld_m will go
# into the L2 halo of fld_f. However, it is a continuous field
# so only the L1 halo will actually be clean.
expected = (" CALL fld_f_proxy%set_dirty()\n"
" CALL fld_f_proxy%set_clean(1)\n"
" !\n"
" CALL fld_f_proxy%halo_exchange(depth=2)\n"
" !\n"
" DO cell=1,mesh_fld_m%get_last_halo_cell(1)\n"
" !\n"
" CALL restrict_test_kernel_code")
assert expected in output
# For the final restriction we need the L2 halo of fld_m to be
# clean. There's no set_clean() call on fld_m because it is
# only updated out to the L1 halo and it is a continuous field
# so the shared dofs in the L1 halo will still be dirty.
expected = (" CALL fld_m_proxy%set_dirty()\n"
" !\n"
" CALL fld_m_proxy%halo_exchange(depth=2)\n"
" !\n"
" DO cell=1,mesh_fld_c%get_last_halo_cell(1)\n"
" !\n"
" CALL restrict_test_kernel_code")
assert expected in output
else:
expected = (
" DO cell=1,fld_c_proxy%vspace%get_ncell()\n"
" !\n"
" CALL prolong_test_kernel_code(nlayers, cell_map_fld_c"
"(:,cell), ncpc_fld_m_fld_c, ncell_fld_m, fld_m_proxy%data, "
"fld_c_proxy%data, ndf_w1, undf_w1, map_w1, undf_w2, "
"map_w2(:,cell))\n"
" END DO \n"
" DO cell=1,fld_m_proxy%vspace%get_ncell()\n"
" !\n"
" CALL prolong_test_kernel_code(nlayers, cell_map_fld_m"
"(:,cell), ncpc_fld_f_fld_m, ncell_fld_f, fld_f_proxy%data, "
"fld_m_proxy%data, ndf_w1, undf_w1, map_w1, undf_w2, "
"map_w2(:,cell))\n"
" END DO \n"
" DO cell=1,fld_m_proxy%vspace%get_ncell()\n"
" !\n"
" CALL restrict_test_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_any_space_1_fld_m, "
"map_any_space_1_fld_m(:,cell), ndf_any_space_2_fld_f, "
"undf_any_space_2_fld_f, map_any_space_2_fld_f)\n"
" END DO \n"
" DO cell=1,fld_c_proxy%vspace%get_ncell()\n"
" !\n"
" CALL restrict_test_kernel_code(nlayers, cell_map_fld_c"
"(:,cell), ncpc_fld_m_fld_c, ncell_fld_m, fld_c_proxy%data, "
"fld_m_proxy%data, undf_any_space_1_fld_c, "
"map_any_space_1_fld_c(:,cell), ndf_any_space_2_fld_m, "
"undf_any_space_2_fld_m, map_any_space_2_fld_m)\n")
assert expected in output
def test_field_restrict(tmpdir, monkeypatch, annexed):
'''Test that we generate correct code for an invoke containing a
single restriction operation (read from fine, write to
coarse). Check when annexed is False and True as we produce a
different number of halo exchanges.
'''
config = Config.get()
dyn_config = config.api_conf("dynamo0.3")
monkeypatch.setattr(dyn_config, "_compute_annexed_dofs", annexed)
_, invoke_info = parse(os.path.join(BASE_PATH,
"22.1_intergrid_restrict.f90"),
api=API)
for distmem in [False, True]:
psy = PSyFactory(API, distributed_memory=distmem).create(invoke_info)
output = str(psy.gen)
print(output)
assert Dynamo0p3Build(tmpdir).code_compiles(psy)
defs = (" USE restrict_test_kernel_mod, "
"ONLY: restrict_test_kernel_code\n"
" USE mesh_map_mod, ONLY: mesh_map_type\n"
" USE mesh_mod, ONLY: mesh_type\n"
" TYPE(field_type), intent(inout) :: field1\n"
" TYPE(field_type), intent(in) :: field2\n")
assert defs in output
defs2 = (
" INTEGER nlayers\n"
" TYPE(field_proxy_type) field1_proxy, field2_proxy\n"
" INTEGER, pointer :: map_any_space_1_field1(:,:) => null(), "
"map_any_space_2_field2(:,:) => null()\n"
" INTEGER ndf_any_space_1_field1, undf_any_space_1_field1, "
"ndf_any_space_2_field2, undf_any_space_2_field2\n"
" INTEGER ncell_field2, ncpc_field2_field1\n"
" INTEGER, pointer :: cell_map_field1(:,:) => null()\n"
" TYPE(mesh_map_type), pointer :: mmap_field2_field1 => "
"null()\n"
" TYPE(mesh_type), pointer :: mesh_field2 => null()\n"
" TYPE(mesh_type), pointer :: mesh_field1 => null()\n")
assert defs2 in output
inits = (" !\n"
" ! Look-up mesh objects and loop limits for inter-grid "
"kernels\n"
" !\n"
" mesh_field2 => field2_proxy%vspace%get_mesh()\n"
" mesh_field1 => field1_proxy%vspace%get_mesh()\n"
" mmap_field2_field1 => mesh_field1%get_mesh_map("
"mesh_field2)\n"
" cell_map_field1 => mmap_field2_field1%"
"get_whole_cell_map()\n")
if distmem:
inits += (" ncell_field2 = mesh_field2%"
"get_last_halo_cell(depth=2)\n")
else:
inits += (" ncell_field2 = field2_proxy%vspace%"
"get_ncell()\n")
inits += (" ncpc_field2_field1 = mmap_field2_field1%"
"get_ntarget_cells_per_source_cell()\n"
" !\n"
" ! Look-up dofmaps for each function space\n"
" !\n"
" map_any_space_1_field1 => field1_proxy%vspace%"
"get_whole_dofmap()\n"
" map_any_space_2_field2 => field2_proxy%vspace%"
"get_whole_dofmap()\n")
assert inits in output
if distmem:
# We write out to the L1 halo on the coarse mesh which means
# we require up-to-date values out to the L2 halo on the fine.
# Since we are incrementing the coarse field we also need
# up-to-date values for it in the L1 halo.
if not annexed:
halo_exchs = (
" ! Call kernels and communication routines\n"
" !\n"
" IF (field1_proxy%is_dirty(depth=1)) THEN\n"
" CALL field1_proxy%halo_exchange(depth=1)\n"
" END IF \n"
" !\n"
" IF (field2_proxy%is_dirty(depth=2)) THEN\n"
" CALL field2_proxy%halo_exchange(depth=2)\n"
" END IF \n"
" !\n"
" DO cell=1,mesh_field1%get_last_halo_cell(1)\n")
else:
halo_exchs = (
" ! Call kernels and communication routines\n"
" !\n"
" IF (field2_proxy%is_dirty(depth=2)) THEN\n"
" CALL field2_proxy%halo_exchange(depth=2)\n"
" END IF \n"
" !\n"
" DO cell=1,mesh_field1%get_last_halo_cell(1)\n")
assert halo_exchs in output
# We pass the whole dofmap for the fine mesh (we are reading from).
# This is associated with the second kernel argument.
kern_call = (
" !\n"
" CALL restrict_test_kernel_code(nlayers, "
"cell_map_field1(:,cell), ncpc_field2_field1, ncell_field2, "
"field1_proxy%data, field2_proxy%data, "
"undf_any_space_1_field1, map_any_space_1_field1(:,cell), "
"ndf_any_space_2_field2, undf_any_space_2_field2, "
"map_any_space_2_field2)\n"
" END DO \n"
" !\n")
assert kern_call in output
if distmem:
set_dirty = " CALL field1_proxy%set_dirty()\n"
assert set_dirty in output
def test_field_prolong(tmpdir):
''' Check that we generate correct psy-layer code for an invoke
containing a kernel that performs a prolongation operation '''
_, invoke_info = parse(os.path.join(BASE_PATH,
"22.0_intergrid_prolong.f90"),
api=API)
for distmem in [False, True]:
psy = PSyFactory(API, distributed_memory=distmem).create(invoke_info)
gen_code = str(psy.gen)
assert Dynamo0p3Build(tmpdir).code_compiles(psy)
expected = (" USE prolong_test_kernel_mod, "
"ONLY: prolong_test_kernel_code\n"
" USE mesh_map_mod, ONLY: mesh_map_type\n"
" USE mesh_mod, ONLY: mesh_type\n"
" TYPE(field_type), intent(inout) :: field1\n"
" TYPE(field_type), intent(in) :: field2\n"
" INTEGER cell\n")
assert expected in gen_code
expected = (
" INTEGER ncell_field1, ncpc_field1_field2\n"
" INTEGER, pointer :: cell_map_field2(:,:) => null()\n"
" TYPE(mesh_map_type), pointer :: "
"mmap_field1_field2 => null()\n"
" TYPE(mesh_type), pointer :: mesh_field2 => null()\n"
" TYPE(mesh_type), pointer :: mesh_field1 => null()\n")
assert expected in gen_code
expected = (
" ! Look-up mesh objects and loop limits for inter-grid "
"kernels\n"
" !\n"
" mesh_field1 => field1_proxy%vspace%get_mesh()\n"
" mesh_field2 => field2_proxy%vspace%get_mesh()\n"
" mmap_field1_field2 => mesh_field2%get_mesh_map"
"(mesh_field1)\n"
" cell_map_field2 => mmap_field1_field2%"
"get_whole_cell_map()\n")
if distmem:
expected += (" ncell_field1 = mesh_field1%get_last_halo_cell("
"depth=2)\n")
else:
expected += \
" ncell_field1 = field1_proxy%vspace%get_ncell()\n"
expected += (" ncpc_field1_field2 = mmap_field1_field2%"
"get_ntarget_cells_per_source_cell()\n")
assert expected in gen_code
if distmem:
# We are writing to a continuous field on the fine mesh, we
# only need to halo swap to depth one on the coarse.
expected = (" IF (field2_proxy%is_dirty(depth=1)) THEN\n"
" CALL field2_proxy%halo_exchange(depth=1)\n"
" END IF \n"
" !\n"
" DO cell=1,mesh_field2%get_last_halo_cell(1)\n")
assert expected in gen_code
else:
assert "DO cell=1,field2_proxy%vspace%get_ncell()\n" in gen_code
expected = (
" CALL prolong_test_kernel_code(nlayers, "
"cell_map_field2(:,cell), ncpc_field1_field2, ncell_field1, "
"field1_proxy%data, field2_proxy%data, ndf_w1, undf_w1, map_w1, "
"undf_w2, map_w2(:,cell))\n"
" END DO \n")
assert expected in gen_code
if distmem:
set_dirty = " CALL field1_proxy%set_dirty()\n"
assert set_dirty in gen_code
def test_gh_inc_nohex_4(tmpdir, monkeypatch):
'''If COMPUTE_ANNEXED_DOFS is False, then a gh_inc access to a field
in a kernel (iterating to the l1 halo) does not require a halo
exchange when the previous writer is known and iterates over cells
to halo(1), halo(2) and halo max depth. Also, if the previous
writer is a gh_inc access and its previous writer is unknown then
it does require a halo exchange if it writes to halo(1) and
requires a speculative halo exchange to halo(n-1) if iterating to
halo(n) and a speculative halo exchange to halo(max_depth-1) if
iterating to the maximum halo depth
'''
# ensure that COMPUTE_ANNEXED_DOFS is False
config = Config.get()
dyn_config = config.api_conf(API)
monkeypatch.setattr(dyn_config, "_compute_annexed_dofs", False)
# parse and get psy schedule
_, info = parse(os.path.join(BASE_PATH,
"14.13_halo_inc_to_inc.f90"),
api=API)
psy = PSyFactory(API, distributed_memory=True).create(info)
schedule = psy.invokes.invoke_list[0].schedule
def check(schedule, f1depth, f2depth):
'''check that the schedule is modified in the expected way. In
particular, check that the depth of the halo exchange for
field 'f1' is what we are expecting
:param schedule: a dynamo0.3 API schedule object
:type schedule: :py:class:`psyclone.dynamo0p3.DynInvokeSchedule`.
:param int f1depth: The expected depth of the halo exchange \
associated with field f1
:param int f2depth: The expected depth of the halo exchange \
associated with field f2
'''
assert len(schedule.children) == 4
haloex1 = schedule.children[0]
haloex2 = schedule.children[1]
loop1 = schedule.children[2]
loop2 = schedule.children[3]
assert isinstance(haloex1, DynHaloExchange)
assert haloex1.field.name == "f1"
assert haloex1._compute_halo_depth() == f1depth
assert haloex1.required() == (True, False)
assert isinstance(haloex2, DynHaloExchange)
assert haloex2.field.name == "f2"
assert haloex2._compute_halo_depth() == f2depth
assert haloex2.required() == (True, False)
assert isinstance(loop1, DynLoop)
assert isinstance(loop2, DynLoop)
# we should now have a speculative halo exchange at the start of
# the schedule for "f1" to depth 1 and "f2" to depth 1
check(schedule, f1depth="1", f2depth="1")
# just check compilation here (not later in this test) as
# compilation of redundant computation is checked separately
assert Dynamo0p3Build(tmpdir).code_compiles(psy)
# make 1st loop iterate over cells to the level 2 halo and check output
rc_trans = Dynamo0p3RedundantComputationTrans()
rc_trans.apply(schedule.children[2], depth=2)
# we should now have a speculative halo exchange at the start of
# the schedule for "f1" to depth 1 and "f2" to depth 2
check(schedule, f1depth="1", f2depth="2")
# make 1st loop iterate over cells to the maximum halo depth and
# check output
rc_trans.apply(schedule.children[2])
# we should now have a speculative halo exchange at the start of
# the schedule for "f1" to depth max halo - 1 and "f2" to max halo
check(schedule, f1depth="mesh%get_halo_depth()-1",
f2depth="mesh%get_halo_depth()")
def test_gh_inc_max(tmpdir, monkeypatch, annexed):
'''Check we generate correct halo exchange bounds when we have
multiple read dependencies. In this case we have a gh_inc with a
read-only reader and a gh_inc reader. We also test when annexed
is False and True as it affects how many halo exchanges are
generated.
'''
config = Config.get()
dyn_config = config.api_conf(API)
monkeypatch.setattr(dyn_config, "_compute_annexed_dofs", annexed)
# parse and get psy schedule
_, info = parse(os.path.join(BASE_PATH,
"14.14_halo_inc_times3.f90"),
api=API)
psy = PSyFactory(API, distributed_memory=True).create(info)
schedule = psy.invokes.invoke_list[0].schedule
rc_trans = Dynamo0p3RedundantComputationTrans()
def check(haloex, depth):
'''check the halo exchange has the expected properties
:param haloex: a dynamo0.3 API halo-exchange object
:type haloex: :py:class:`psyclone.dynamo0p3.DynHaloExchange`.
:param int depth: The expected depth of the halo exchange \
passed in as the first argument
'''
assert isinstance(haloex, DynHaloExchange)
assert haloex.field.name == "f1"
assert haloex.required() == (True, True)
assert haloex._compute_halo_depth() == depth
if annexed:
haloidx = 2
loop1idx = 3
loop2idx = 5
else:
haloidx = 4
loop1idx = 5
loop2idx = 7
# f1 halo exchange should be depth 1 : max(1,0)
haloex = schedule.children[haloidx]
check(haloex, "1")
rc_trans.apply(schedule.children[loop2idx], depth=2)
# f1 halo exchange should still be depth 1 : max(1,1)
haloex = schedule.children[haloidx]
check(haloex, "1")
rc_trans.apply(schedule.children[loop2idx], depth=3)
# f1 halo exchange should be depth 2 (max(1,2)
haloex = schedule.children[haloidx]
check(haloex, "2")
rc_trans.apply(schedule.children[loop2idx])
# f1 halo exchange should be depth max(1,max-1)
haloex = schedule.children[haloidx]
check(haloex, "max(mesh%get_halo_depth()-1,1)")
# just check compilation here as it is the most
# complicated. (Note, compilation of redundant computation is
# checked separately)
assert Dynamo0p3Build(tmpdir).code_compiles(psy)
rc_trans.apply(schedule.children[loop1idx])
# f1 halo exchange should be depth max
haloex = schedule.children[haloidx]
check(haloex, "mesh%get_halo_depth()")
def test_field_xyoz(tmpdir):
''' Tests that a call, with a set of fields requiring XYoZ
quadrature, produces correct code. '''
_, invoke_info = parse(os.path.join(BASE_PATH,
"1.1.0_single_invoke_xyoz_qr.f90"),
api=API)
psy = PSyFactory(API, distributed_memory=True).create(invoke_info)
generated_code = str(psy.gen)
print(generated_code)
assert Dynamo0p3Build(tmpdir).code_compiles(psy)
output_decls = (
" SUBROUTINE invoke_0_testkern_qr_type(f1, f2, m1, a, m2, istp,"
" qr)\n"
" USE testkern_qr, ONLY: testkern_qr_code\n"
" USE quadrature_xyoz_mod, ONLY: quadrature_xyoz_type, "
"quadrature_xyoz_proxy_type\n"
" USE function_space_mod, ONLY: BASIS, DIFF_BASIS\n"
" USE mesh_mod, ONLY: mesh_type\n"
" REAL(KIND=r_def), intent(in) :: a\n"
" INTEGER, intent(in) :: istp\n"
" TYPE(field_type), intent(inout) :: f1\n"
" TYPE(field_type), intent(in) :: f2, m1, m2\n"
" TYPE(quadrature_xyoz_type), intent(in) :: qr\n"
" INTEGER cell\n"
" REAL(KIND=r_def), allocatable :: basis_w1_qr(:,:,:,:), "
"diff_basis_w2_qr(:,:,:,:), basis_w3_qr(:,:,:,:), "
"diff_basis_w3_qr(:,:,:,:)\n"
" INTEGER dim_w1, diff_dim_w2, dim_w3, diff_dim_w3\n"
" REAL(KIND=r_def), pointer :: weights_xy_qr(:) => null(), "
"weights_z_qr(:) => null()\n"
" INTEGER np_xy_qr, np_z_qr\n"
" INTEGER nlayers\n"
" TYPE(field_proxy_type) f1_proxy, f2_proxy, m1_proxy, m2_proxy\n"
" TYPE(quadrature_xyoz_proxy_type) qr_proxy\n"
" INTEGER, pointer :: map_w1(:,:) => null(), "
"map_w2(:,:) => null(), map_w3(:,:) => null()\n"
" INTEGER ndf_w1, undf_w1, ndf_w2, undf_w2, ndf_w3, undf_w3\n")
assert output_decls in generated_code
init_output = (" !\n"
" ! Initialise field and/or operator proxies\n"
" !\n"
" f1_proxy = f1%get_proxy()\n"
" f2_proxy = f2%get_proxy()\n"
" m1_proxy = m1%get_proxy()\n"
" m2_proxy = m2%get_proxy()\n"
" !\n"
" ! Initialise number of layers\n"
" !\n"
" nlayers = f1_proxy%vspace%get_nlayers()\n"
" !\n"
" ! Create a mesh object\n"
" !\n"
" mesh => f1_proxy%vspace%get_mesh()\n"
" !\n"
" ! Look-up dofmaps for each function space\n"
" !\n"
" map_w1 => f1_proxy%vspace%get_whole_dofmap()\n"
" map_w2 => f2_proxy%vspace%get_whole_dofmap()\n"
" map_w3 => m2_proxy%vspace%get_whole_dofmap()\n"
" !\n"
" ! Initialise number of DoFs for w1\n"
" !\n"
" ndf_w1 = f1_proxy%vspace%get_ndf()\n"
" undf_w1 = f1_proxy%vspace%get_undf()\n"
" !\n"
" ! Initialise number of DoFs for w2\n"
" !\n"
" ndf_w2 = f2_proxy%vspace%get_ndf()\n"
" undf_w2 = f2_proxy%vspace%get_undf()\n"
" !\n"
" ! Initialise number of DoFs for w3\n"
" !\n"
" ndf_w3 = m2_proxy%vspace%get_ndf()\n"
" undf_w3 = m2_proxy%vspace%get_undf()\n"
" !\n"
" ! Look-up quadrature variables\n"
" !\n"
" qr_proxy = qr%get_quadrature_proxy()\n"
" np_xy_qr = qr_proxy%np_xy\n"
" np_z_qr = qr_proxy%np_z\n"
" weights_xy_qr => qr_proxy%weights_xy\n"
" weights_z_qr => qr_proxy%weights_z\n")
assert init_output in generated_code
compute_output = (
" !\n"
" ! Allocate basis/diff-basis arrays\n"
" !\n"
" dim_w1 = f1_proxy%vspace%get_dim_space()\n"
" diff_dim_w2 = f2_proxy%vspace%get_dim_space_diff()\n"
" dim_w3 = m2_proxy%vspace%get_dim_space()\n"
" diff_dim_w3 = m2_proxy%vspace%get_dim_space_diff()\n"
" ALLOCATE (basis_w1_qr(dim_w1, ndf_w1, np_xy_qr, np_z_qr))\n"
" ALLOCATE (diff_basis_w2_qr(diff_dim_w2, ndf_w2, np_xy_qr, "
"np_z_qr))\n"
" ALLOCATE (basis_w3_qr(dim_w3, ndf_w3, np_xy_qr, np_z_qr))\n"
" ALLOCATE (diff_basis_w3_qr(diff_dim_w3, ndf_w3, np_xy_qr, "
"np_z_qr))\n"
" !\n"
" ! Compute basis/diff-basis arrays\n"
" !\n"
" CALL qr%compute_function(BASIS, f1_proxy%vspace, dim_w1, "
"ndf_w1, basis_w1_qr)\n"
" CALL qr%compute_function(DIFF_BASIS, f2_proxy%vspace, "
"diff_dim_w2, ndf_w2, diff_basis_w2_qr)\n"
" CALL qr%compute_function(BASIS, m2_proxy%vspace, dim_w3, "
"ndf_w3, basis_w3_qr)\n"
" CALL qr%compute_function(DIFF_BASIS, m2_proxy%vspace, "
"diff_dim_w3, ndf_w3, diff_basis_w3_qr)\n"
" !\n"
" ! Call kernels and communication routines\n"
" !\n"
" IF (f2_proxy%is_dirty(depth=1)) THEN\n"
" CALL f2_proxy%halo_exchange(depth=1)\n"
" END IF \n"
" !\n"
" IF (m1_proxy%is_dirty(depth=1)) THEN\n"
" CALL m1_proxy%halo_exchange(depth=1)\n"
" END IF \n"
" !\n"
" IF (m2_proxy%is_dirty(depth=1)) THEN\n"
" CALL m2_proxy%halo_exchange(depth=1)\n"
" END IF \n"
" !\n"
" DO cell=1,mesh%get_last_halo_cell(1)\n"
" !\n"
" CALL testkern_qr_code(nlayers, f1_proxy%data, f2_proxy%data, "
"m1_proxy%data, a, m2_proxy%data, istp, ndf_w1, undf_w1, "
"map_w1(:,cell), basis_w1_qr, ndf_w2, undf_w2, map_w2(:,cell), "
"diff_basis_w2_qr, ndf_w3, undf_w3, map_w3(:,cell), basis_w3_qr, "
"diff_basis_w3_qr, np_xy_qr, np_z_qr, weights_xy_qr, weights_z_qr)\n"
" END DO \n"
" !\n"
" ! Set halos dirty/clean for fields modified in the above loop\n"
" !\n"
" CALL f1_proxy%set_dirty()\n"
" !\n"
" !\n"
" ! Deallocate basis arrays\n"
" !\n"
" DEALLOCATE (basis_w1_qr, basis_w3_qr, diff_basis_w2_qr, "
"diff_basis_w3_qr)\n"
" !\n"
" END SUBROUTINE invoke_0_testkern_qr_type")
assert compute_output in generated_code
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)
def test_operator_different_spaces(tmpdir):
''' Tests that an operator with different to and from spaces is
implemented correctly in the PSy layer '''
_, invoke_info = parse(os.path.join(BASE_PATH,
"10.3_operator_different_spaces.f90"),
api=TEST_API)
psy = PSyFactory(TEST_API, distributed_memory=True).create(invoke_info)
generated_code = str(psy.gen)
assert Dynamo0p3Build(tmpdir).code_compiles(psy)
decl_output = (
" SUBROUTINE invoke_0_assemble_weak_derivative_w3_w2_kernel_type"
"(mapping, chi, qr)\n"
" USE assemble_weak_derivative_w3_w2_kernel_mod, ONLY: "
"assemble_weak_derivative_w3_w2_kernel_code\n"
" USE quadrature_xyoz_mod, ONLY: quadrature_xyoz_type, "
"quadrature_xyoz_proxy_type\n"
" USE function_space_mod, ONLY: BASIS, DIFF_BASIS\n"
" USE mesh_mod, ONLY: mesh_type\n"
" TYPE(field_type), intent(in) :: chi(3)\n"
" TYPE(operator_type), intent(inout) :: mapping\n"
" TYPE(quadrature_xyoz_type), intent(in) :: qr\n"
" INTEGER cell\n"
" REAL(KIND=r_def), allocatable :: diff_basis_w0_qr(:,:,:,:), "
"basis_w3_qr(:,:,:,:), diff_basis_w2_qr(:,:,:,:)\n"
" INTEGER diff_dim_w0, dim_w3, diff_dim_w2\n"
" REAL(KIND=r_def), pointer :: weights_xy_qr(:) => null(), "
"weights_z_qr(:) => null()\n"
" INTEGER np_xy_qr, np_z_qr\n"
" INTEGER nlayers\n"
" TYPE(operator_proxy_type) mapping_proxy\n"
" TYPE(field_proxy_type) chi_proxy(3)\n"
" TYPE(quadrature_xyoz_proxy_type) qr_proxy\n"
" INTEGER, pointer :: map_w0(:,:) => null()\n"
" INTEGER ndf_w3, ndf_w2, ndf_w0, undf_w0\n"
" TYPE(mesh_type), pointer :: mesh => null()\n"
" INTEGER, pointer :: orientation_w2(:) => null()\n")
assert decl_output in generated_code
output = (
" !\n"
" ! Initialise field and/or operator proxies\n"
" !\n"
" mapping_proxy = mapping%get_proxy()\n"
" chi_proxy(1) = chi(1)%get_proxy()\n"
" chi_proxy(2) = chi(2)%get_proxy()\n"
" chi_proxy(3) = chi(3)%get_proxy()\n"
" !\n"
" ! Initialise number of layers\n"
" !\n"
" nlayers = mapping_proxy%fs_from%get_nlayers()\n"
" !\n"
" ! Create a mesh object\n"
" !\n"
" mesh => mapping_proxy%fs_from%get_mesh()\n"
" !\n"
" ! Look-up dofmaps for each function space\n"
" !\n"
" map_w0 => chi_proxy(1)%vspace%get_whole_dofmap()\n"
" !\n"
" ! Initialise number of DoFs for w3\n"
" !\n"
" ndf_w3 = mapping_proxy%fs_to%get_ndf()\n"
" !\n"
" ! Initialise number of DoFs for w2\n"
" !\n"
" ndf_w2 = mapping_proxy%fs_from%get_ndf()\n"
" !\n"
" ! Initialise number of DoFs for w0\n"
" !\n"
" ndf_w0 = chi_proxy(1)%vspace%get_ndf()\n"
" undf_w0 = chi_proxy(1)%vspace%get_undf()\n"
" !\n"
" ! Look-up quadrature variables\n"
" !\n"
" qr_proxy = qr%get_quadrature_proxy()\n"
" np_xy_qr = qr_proxy%np_xy\n"
" np_z_qr = qr_proxy%np_z\n"
" weights_xy_qr => qr_proxy%weights_xy\n"
" weights_z_qr => qr_proxy%weights_z\n"
" !\n"
" ! Allocate basis/diff-basis arrays\n"
" !\n"
" diff_dim_w0 = chi_proxy(1)%vspace%get_dim_space_diff()\n"
" dim_w3 = mapping_proxy%fs_to%get_dim_space()\n"
" diff_dim_w2 = mapping_proxy%fs_from%get_dim_space_diff()\n"
" ALLOCATE (diff_basis_w0_qr(diff_dim_w0, ndf_w0, np_xy_qr, "
"np_z_qr))\n"
" ALLOCATE (basis_w3_qr(dim_w3, ndf_w3, np_xy_qr, np_z_qr))\n"
" ALLOCATE (diff_basis_w2_qr(diff_dim_w2, ndf_w2, np_xy_qr, "
"np_z_qr))\n"
" !\n"
" ! Compute basis/diff-basis arrays\n"
" !\n"
" CALL qr%compute_function(DIFF_BASIS, chi_proxy(1)%vspace, "
"diff_dim_w0, ndf_w0, diff_basis_w0_qr)\n"
" CALL qr%compute_function(BASIS, mapping_proxy%fs_to, "
"dim_w3, ndf_w3, basis_w3_qr)\n"
" CALL qr%compute_function(DIFF_BASIS, mapping_proxy%fs_from, "
"diff_dim_w2, ndf_w2, diff_basis_w2_qr)\n"
" !\n"
" ! Call kernels and communication routines\n"
" !\n"
" IF (chi_proxy(1)%is_dirty(depth=1)) THEN\n"
" CALL chi_proxy(1)%halo_exchange(depth=1)\n"
" END IF \n"
" !\n"
" IF (chi_proxy(2)%is_dirty(depth=1)) THEN\n"
" CALL chi_proxy(2)%halo_exchange(depth=1)\n"
" END IF \n"
" !\n"
" IF (chi_proxy(3)%is_dirty(depth=1)) THEN\n"
" CALL chi_proxy(3)%halo_exchange(depth=1)\n"
" END IF \n"
" !\n"
" DO cell=1,mesh%get_last_halo_cell(1)\n"
" !\n"
" orientation_w2 => mapping_proxy%fs_from%get_cell_orientation("
"cell)\n"
" !\n"
" CALL assemble_weak_derivative_w3_w2_kernel_code(cell, "
"nlayers, mapping_proxy%ncell_3d, mapping_proxy%local_stencil, "
"chi_proxy(1)%data, chi_proxy(2)%data, chi_proxy(3)%data, ndf_w3, "
"basis_w3_qr, ndf_w2, diff_basis_w2_qr, orientation_w2, "
"ndf_w0, undf_w0, map_w0(:,cell), diff_basis_w0_qr, "
"np_xy_qr, np_z_qr, weights_xy_qr, weights_z_qr)\n"
" END DO \n"
" !\n"
" ! Deallocate basis arrays\n"
" !\n"
" DEALLOCATE (basis_w3_qr, diff_basis_w0_qr, diff_basis_w2_qr)\n"
" !\n"
" END SUBROUTINE invoke_0_assemble_weak_derivative_w3_w2_kernel_"
"type")
assert output in generated_code