diff --git a/examples/lfric/code/testkern_mod.F90 b/examples/lfric/code/testkern_mod.F90 index 332175682a..e3df0a95bb 100644 --- a/examples/lfric/code/testkern_mod.F90 +++ b/examples/lfric/code/testkern_mod.F90 @@ -30,7 +30,7 @@ ! OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. ! ----------------------------------------------------------------------------- ! Authors: R. W. Ford and A. R. Porter, STFC Daresbury Laboratory -! Modified: I. Kavcic, Met Office +! Modified: I. Kavcic and A. Pirrie, Met Office module testkern_mod @@ -42,12 +42,13 @@ module testkern_mod implicit none type, extends(kernel_type) :: testkern_type - type(arg_type), dimension(5) :: meta_args = & - (/ arg_type(gh_scalar, gh_real, gh_read), & - arg_type(gh_field, gh_real, gh_inc, w1), & - arg_type(gh_field, gh_real, gh_read, w2), & - arg_type(gh_field, gh_real, gh_read, w2), & - arg_type(gh_field, gh_real, gh_read, w3) & + type(arg_type), dimension(6) :: meta_args = & + (/ arg_type(gh_scalar, gh_real, gh_read ), & + arg_type(gh_scalar_array, gh_real, gh_read, 2 ), & + arg_type(gh_field, gh_real, gh_inc, w1), & + arg_type(gh_field, gh_real, gh_read, w2), & + arg_type(gh_field, gh_real, gh_read, w2), & + arg_type(gh_field, gh_real, gh_read, w3) & /) integer :: operates_on = cell_column contains @@ -56,10 +57,11 @@ module testkern_mod contains - subroutine testkern_code(nlayers, ascalar, & - fld1, fld2, fld3, fld4, & - ndf_w1, undf_w1, map_w1, & - ndf_w2, undf_w2, map_w2, & + subroutine testkern_code(nlayers, ascalar, & + dims_ascalar_array, ascalar_array, & + fld1, fld2, fld3, fld4, & + ndf_w1, undf_w1, map_w1, & + ndf_w2, undf_w2, map_w2, & ndf_w3, undf_w3, map_w3) implicit none @@ -71,6 +73,8 @@ subroutine testkern_code(nlayers, ascalar, & integer(kind=i_def), intent(in), dimension(ndf_w1) :: map_w1 integer(kind=i_def), intent(in), dimension(ndf_w2) :: map_w2 integer(kind=i_def), intent(in), dimension(ndf_w3) :: map_w3 + integer(kind=i_def), intent(in), dimension(2) :: dims_ascalar_array + real(kind=r_def), intent(in), dimension(dims_ascalar_array(1), dims_ascalar_array(2)) :: ascalar_array real(kind=r_def), intent(in) :: ascalar real(kind=r_def), intent(inout), dimension(undf_w1) :: fld1 real(kind=r_def), intent(in), dimension(undf_w2) :: fld2 diff --git a/examples/lfric/eg1/single_invoke.x90 b/examples/lfric/eg1/single_invoke.x90 index df91b42e43..a046ad2c40 100644 --- a/examples/lfric/eg1/single_invoke.x90 +++ b/examples/lfric/eg1/single_invoke.x90 @@ -32,7 +32,7 @@ ! POSSIBILITY OF SUCH DAMAGE. ! ----------------------------------------------------------------------------- ! Authors: A. R. Porter and R. W. Ford, STFC Daresbury Laboratory -! I. Kavcic, Met Office +! I. Kavcic and A. Pirrie, Met Office program single_invoke @@ -44,10 +44,11 @@ program single_invoke implicit none - type(field_type) :: f1, f2, m1, m2 - real(r_def) :: a + type(field_type) :: f1, f2, m1, m2 + real(r_def) :: a + real(r_def), dimension(50,100) :: ascalar_array call invoke( setval_c(f1, 0.0_r_def), & - testkern_type(a, f1, f2, m1, m2) ) + testkern_type(a, ascalar_array, f1, f2, m1, m2) ) end program single_invoke diff --git a/examples/lfric/eg2/multi_invoke_mod.x90 b/examples/lfric/eg2/multi_invoke_mod.x90 index 0a5f9ff3dc..963557266b 100644 --- a/examples/lfric/eg2/multi_invoke_mod.x90 +++ b/examples/lfric/eg2/multi_invoke_mod.x90 @@ -32,7 +32,7 @@ ! POSSIBILITY OF SUCH DAMAGE. ! ----------------------------------------------------------------------------- ! Author R. W. Ford and A. R. Porter, STFC Daresbury Laboratory -! Modified I. Kavcic, Met Office +! Modified I. Kavcic and A. Pirrie, Met Office program multi_invoke @@ -45,8 +45,9 @@ program multi_invoke implicit none - type(field_type) :: f1, f2, m1, m2 - real(r_def) :: a + type(field_type) :: f1, f2, m1, m2 + real(r_def) :: a + real(r_def), dimension(50,100) :: ascalar_array call invoke( setval_c(f1, 0.0_r_def), & setval_c(f2, 0.0_r_def), & @@ -54,7 +55,7 @@ program multi_invoke setval_c(m2, 0.0_r_def) ) ! This example artificially has two separate 'invoke()' calls. In practice ! these would be merged into a single call. - call invoke( testkern_type(a, f1, f2, m1, m2), & - testkern_type(a, f1, f2, m1, m2) ) + call invoke( testkern_type(a, ascalar_array, f1, f2, m1, m2), & + testkern_type(a, ascalar_array, f1, f2, m1, m2) ) end program multi_invoke diff --git a/src/psyclone/domain/lfric/transformations/lfric_loop_fuse_trans.py b/src/psyclone/domain/lfric/transformations/lfric_loop_fuse_trans.py index 488e64f2ee..98db3a8d96 100644 --- a/src/psyclone/domain/lfric/transformations/lfric_loop_fuse_trans.py +++ b/src/psyclone/domain/lfric/transformations/lfric_loop_fuse_trans.py @@ -168,7 +168,6 @@ def validate(self, node1: LFRicLoop, node2: LFRicLoop, # 2.2) If 'same_space' is true check that both function spaces are # the same or that at least one of the nodes is on ANY_SPACE. The # former case is convenient when loop fusion is applied generically. - if same_space: if node1_fs_name == node2_fs_name: pass @@ -182,12 +181,15 @@ def validate(self, node1: LFRicLoop, node2: LFRicLoop, # 2.3.1) Check whether one or more of the function spaces # is ANY_SPACE without the 'same_space' flag if node_on_any_space: - raise TransformationError( - f"Error in {self.name} transformation: One or more of the " - f"iteration spaces is unknown ('ANY_SPACE') so loop fusion" - f" might be invalid. If you know the spaces are the same " - f"then please set the 'same_space' optional argument to " - f"'True'.") + # If the nodes are on ANY_SPACE, but those are the same + # space, we can fuse. + if node1_fs_name != node2_fs_name: + raise TransformationError( + f"Error in {self.name} transformation: One or more of " + f"the iteration spaces is unknown ('ANY_SPACE') so " + f"loop fusion might be invalid. If you know the " + f"spaces are the same then please set the " + f"'same_space' optional argument to 'True'.") # 2.3.2) Check whether specific function spaces are the # same. If they are not, the loop fusion is still possible # but only when both function spaces are discontinuous diff --git a/src/psyclone/tests/domain/lfric/transformations/lfric_transformations_test.py b/src/psyclone/tests/domain/lfric/transformations/lfric_transformations_test.py index 1619d1eac7..e13ec177e8 100644 --- a/src/psyclone/tests/domain/lfric/transformations/lfric_transformations_test.py +++ b/src/psyclone/tests/domain/lfric/transformations/lfric_transformations_test.py @@ -3204,18 +3204,53 @@ def test_multi_builtins_fuse_error(): "reduction") in str(excinfo.value) -def test_loop_fuse_error(dist_mem): - '''Test that we raise an exception in loop fusion if one or more of - the loops has an any_space iteration space.''' - _, invoke = get_invoke("15.14.2_multiple_set_kernels.f90", - TEST_API, idx=0, dist_mem=dist_mem) +def test_loop_fuse_any_space(tmpdir, dist_mem): + '''Test that we correctly fuse two or more of loops that are on + any_space iteration space.''' + psy, invoke = get_invoke("15.14.2_multiple_set_kernels.f90", + TEST_API, idx=0, dist_mem=dist_mem) schedule = invoke.schedule ftrans = LFRicLoopFuseTrans() - with pytest.raises(TransformationError) as excinfo: - ftrans.apply(schedule.children[0], schedule.children[1]) - assert ("One or more of the iteration spaces is unknown " - "('ANY_SPACE') so loop fusion might be " - "invalid") in str(excinfo.value) + + # Fuses the first two loops + ftrans.apply(schedule.children[0], schedule.children[1]) + code = str(psy.gen) + assert ( + "do df = loop0_start, loop0_stop, 1\n" + " ! Built-in: setval_c (set a real-valued field to a real scalar " + "value)\n" + " f1_data(df) = fred\n" + "\n" + " ! Built-in: setval_c (set a real-valued field to a real scalar " + "value)\n" + " f2_data(df) = 3.0_r_def\n" + " enddo\n") in code + assert ( + "do df = loop1_start, loop1_stop, 1\n" + " ! Built-in: setval_c (set a real-valued field to a real scalar " + "value)\n" + " f3_data(df) = ginger\n" + " enddo\n") in code + + # Fuses the combined loop with the third loop + ftrans.apply(schedule.children[0], schedule.children[1]) + code = str(psy.gen) + assert ( + "do df = loop0_start, loop0_stop, 1\n" + " ! Built-in: setval_c (set a real-valued field to a real scalar " + "value)\n" + " f1_data(df) = fred\n" + "\n" + " ! Built-in: setval_c (set a real-valued field to a real scalar " + "value)\n" + " f2_data(df) = 3.0_r_def\n" + "\n" + " ! Built-in: setval_c (set a real-valued field to a real scalar " + "value)\n" + " f3_data(df) = ginger\n" + " enddo\n") in code + + assert LFRicBuild(tmpdir).code_compiles(psy) # Repeat the reduction tests for the reproducible version