def test_prolong_vector(tmpdir, f90, f90flags):
''' 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).create(invoke_info)
output = str(psy.gen)
print output
if utils.TEST_COMPILE:
assert utils.code_compiles(API, psy, tmpdir, f90, f90flags)
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_field_qr_deref(tmpdir, f90, f90flags):
''' 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)
if utils.TEST_COMPILE:
assert utils.code_compiles(API, psy, tmpdir, f90, f90flags)
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 code_compiles_invalid_api(tmpdir, f90, f90flags):
''' Check that utils.code_compiles() reject an unrecognised API '''
with pytest.raises(utils.CompileError) as excinfo:
utils.code_compiles("not_an_api", "fake_psy", tmpdir, f90, f90flags)
assert "Unsupported API in " in str(excinfo)
def test_field_xyoz(tmpdir, f90, f90flags):
''' 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).create(invoke_info)
generated_code = str(psy.gen)
print generated_code
if utils.TEST_COMPILE:
assert utils.code_compiles(API, psy, tmpdir, f90, f90flags)
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(:,:,:,:), "
"basis_w3_qr(:,:,:,:), diff_basis_w2_qr(:,:,:,:), "
"diff_basis_w3_qr(:,:,:,:)\n"
" INTEGER dim_w1, dim_w3, diff_dim_w2, 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 ndf_w1, undf_w1, ndf_w2, undf_w2, ndf_w3, undf_w3\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_w2(:,:) => null(), "
"map_w3(:,:) => null(), map_w1(:,:) => null()\n"
" TYPE(mesh_type), pointer :: mesh => null()\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%get_mesh()\n"
" !\n"
" ! Look-up dofmaps for each function space\n"
" !\n"
" map_w2 => f2_proxy%vspace%get_whole_dofmap()\n"
" map_w3 => m2_proxy%vspace%get_whole_dofmap()\n"
" map_w1 => f1_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 arrays\n"
" !\n"
" dim_w1 = f1_proxy%vspace%get_dim_space()\n"
" ALLOCATE (basis_w1_qr(dim_w1, ndf_w1, np_xy_qr, np_z_qr))\n"
" dim_w3 = m2_proxy%vspace%get_dim_space()\n"
" ALLOCATE (basis_w3_qr(dim_w3, ndf_w3, np_xy_qr, np_z_qr))\n"
" !\n"
" ! Allocate differential basis arrays\n"
" !\n"
" diff_dim_w2 = f2_proxy%vspace%get_dim_space_diff()\n"
" ALLOCATE (diff_basis_w2_qr(diff_dim_w2, ndf_w2, np_xy_qr, "
"np_z_qr))\n"
" diff_dim_w3 = m2_proxy%vspace%get_dim_space_diff()\n"
" ALLOCATE (diff_basis_w3_qr(diff_dim_w3, ndf_w3, np_xy_qr, "
"np_z_qr))\n"
" !\n"
" ! Compute basis arrays\n"
" !\n"
" CALL qr%compute_function(BASIS, f1_proxy%vspace, dim_w1, "
"ndf_w1, basis_w1_qr)\n"
" CALL qr%compute_function(BASIS, m2_proxy%vspace, dim_w3, "
"ndf_w3, basis_w3_qr)\n"
" !\n"
" ! Compute differential basis arrays\n"
" !\n"
" CALL qr%compute_function(DIFF_BASIS, f2_proxy%vspace, "
"diff_dim_w2, ndf_w2, diff_basis_w2_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 (diff_basis_w2_qr, basis_w1_qr, basis_w3_qr, "
"diff_basis_w3_qr)\n"
" !\n"
" END SUBROUTINE invoke_0_testkern_qr_type")
assert compute_output in generated_code
def test_restrict_prolong_chain(tmpdir, f90, f90flags):
''' 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)
for distmem in [False, True]:
psy = PSyFactory(API, distributed_memory=distmem).create(invoke_info)
output = str(psy.gen)
if utils.TEST_COMPILE:
assert utils.code_compiles(API, psy, tmpdir, f90, f90flags)
expected = (
" ! Look-up mesh objects and loop limits for inter-grid "
"kernels\n"
" !\n"
" mesh_fld_m => fld_m%get_mesh()\n"
" mesh_fld_c => fld_c%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")
assert expected in output
if distmem:
expected = (
" 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"
" mesh_fld_f => fld_f%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"
" 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")
else:
expected = (
" 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"
" mesh_fld_f => fld_f%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"
" 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")
assert expected in output
# Check that we haven't got duplicated output
assert output.count("mesh_fld_m => fld_m%get_mesh") == 1
assert output.count("ncell_fld_m = ") == 1
assert output.count("ncell_fld_f = ") == 1
if distmem:
# Have a potential halo exchange before 1st prolong
expected = (
" 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
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"
" 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_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_kernel_code")
assert expected in output
else:
expected = (
" DO cell=1,fld_c_proxy%vspace%get_ncell()\n"
" !\n"
" CALL prolong_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_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_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_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, f90, f90flags):
''' Test that we generate correct code for an invoke containing a
single restriction operation (read from find, write to coarse) '''
_, 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
if utils.TEST_COMPILE:
assert utils.code_compiles(API, psy, tmpdir, f90, f90flags)
defs = (
" USE restrict_kernel_mod, ONLY: restrict_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_2_field2(:,:) => null(), "
"map_any_space_1_field1(:,:) => null()\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%get_mesh()\n"
" mesh_field1 => field1%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_2_field2 => field2_proxy%vspace%"
"get_whole_dofmap()\n"
" map_any_space_1_field1 => field1_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.
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")
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_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, f90, f90flags):
''' 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)
if utils.TEST_COMPILE:
assert utils.code_compiles(API, psy, tmpdir, f90, f90flags)
expected = (
" USE prolong_kernel_mod, ONLY: prolong_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%get_mesh()\n"
" mesh_field2 => field2%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_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