From 8ff0cc7acbe4df2b15e30010f98a483790b11780 Mon Sep 17 00:00:00 2001 From: Ahmad Nawab Date: Mon, 2 Oct 2023 16:33:42 +0200 Subject: [PATCH 1/4] Add tests to check correct placement of loop-bound assignment --- .../tests/test_single_column_coalesced.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/transformations/tests/test_single_column_coalesced.py b/transformations/tests/test_single_column_coalesced.py index 619a38f52..62c4b4d26 100644 --- a/transformations/tests/test_single_column_coalesced.py +++ b/transformations/tests/test_single_column_coalesced.py @@ -567,6 +567,13 @@ def test_scc_hoist_multiple_kernels_loops(frontend, horizontal, vertical, blocki assert driver_loops[2].variable == 'jk' assert driver_loops[2].bounds == '2:nz' + # check location of loop-bound assignment + assign = FindNodes(Assignment).visit(driver_loops[0])[0] + assert assign.lhs == 'end' + assert assign.rhs == 'nlon-nb' + assigns = FindNodes(Assignment).visit(driver_loops[1]) + assert not assign in assigns + assert driver_loops[4] in FindNodes(Loop).visit(driver_loops[3].body) assert driver_loops[5] in FindNodes(Loop).visit(driver_loops[3].body) assert driver_loops[6] in FindNodes(Loop).visit(driver_loops[3].body) @@ -596,6 +603,13 @@ def test_scc_hoist_multiple_kernels_loops(frontend, horizontal, vertical, blocki assert driver_loops[10].variable == 'jk' assert driver_loops[10].bounds == '2:nz' + # check location of loop-bound assignment + assign = FindNodes(Assignment).visit(driver_loops[8])[0] + assert assign.lhs == 'end' + assert assign.rhs == 'nlon-nb' + assigns = FindNodes(Assignment).visit(driver_loops[9]) + assert not assign in assigns + rmtree(basedir) From 6e615f07ef2da766f057de09cc850784542819bf Mon Sep 17 00:00:00 2001 From: Ahmad Nawab Date: Mon, 2 Oct 2023 16:43:37 +0200 Subject: [PATCH 2/4] Extend trim_vector_sections functionality to driver loop --- transformations/tests/test_single_column_coalesced.py | 11 +++++++---- .../transformations/single_column_coalesced_vector.py | 2 +- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/transformations/tests/test_single_column_coalesced.py b/transformations/tests/test_single_column_coalesced.py index 62c4b4d26..8294adcb6 100644 --- a/transformations/tests/test_single_column_coalesced.py +++ b/transformations/tests/test_single_column_coalesced.py @@ -410,7 +410,8 @@ def test_scc_hoist_multiple_kernels(frontend, horizontal, vertical, blocking): @pytest.mark.parametrize('frontend', available_frontends()) -def test_scc_hoist_multiple_kernels_loops(frontend, horizontal, vertical, blocking): +@pytest.mark.parametrize('trim_vector_sections', [True, False]) +def test_scc_hoist_multiple_kernels_loops(frontend, trim_vector_sections, horizontal, vertical, blocking): """ Test hoisting of column temporaries to "driver" level. """ @@ -429,7 +430,6 @@ def test_scc_hoist_multiple_kernels_loops(frontend, horizontal, vertical, blocki !$loki driver-loop do b=1, nb end = nlon - nb - !$loki separator do jk = 2, nz do jl = start, end q(jl, jk, b) = 2.0 * jk * jl @@ -523,7 +523,7 @@ def test_scc_hoist_multiple_kernels_loops(frontend, horizontal, vertical, blocki kernel = scheduler.item_map["kernel_mod#kernel"].routine transformation = (SCCBaseTransformation(horizontal=horizontal, directive='openacc'),) - transformation += (SCCDevectorTransformation(horizontal=horizontal, trim_vector_sections=False),) + transformation += (SCCDevectorTransformation(horizontal=horizontal, trim_vector_sections=trim_vector_sections),) transformation += (SCCDemoteTransformation(horizontal=horizontal),) transformation += (SCCRevectorTransformation(horizontal=horizontal),) transformation += (SCCAnnotateTransformation( @@ -572,7 +572,10 @@ def test_scc_hoist_multiple_kernels_loops(frontend, horizontal, vertical, blocki assert assign.lhs == 'end' assert assign.rhs == 'nlon-nb' assigns = FindNodes(Assignment).visit(driver_loops[1]) - assert not assign in assigns + if trim_vector_sections: + assert not assign in assigns + else: + assert assign in assigns assert driver_loops[4] in FindNodes(Loop).visit(driver_loops[3].body) assert driver_loops[5] in FindNodes(Loop).visit(driver_loops[3].body) diff --git a/transformations/transformations/single_column_coalesced_vector.py b/transformations/transformations/single_column_coalesced_vector.py index db99e4d5a..446e7435d 100644 --- a/transformations/transformations/single_column_coalesced_vector.py +++ b/transformations/transformations/single_column_coalesced_vector.py @@ -228,7 +228,7 @@ def process_driver(self, routine, targets=()): new_driver_loop = loop.clone(body=new_driver_loop) sections = self.extract_vector_sections(new_driver_loop.body, self.horizontal) if self.trim_vector_sections: - sections = self.get_trimmed_sections(routine, self.horizontal, sections) + sections = self.get_trimmed_sections(new_driver_loop, self.horizontal, sections) section_mapper = {s: ir.Section(body=s, label='vector_section') for s in sections} new_driver_loop = NestedTransformer(section_mapper).visit(new_driver_loop) driver_loop_map[loop] = new_driver_loop From 186ff678ed161cb98f1276beed7bbc19e309bf0f Mon Sep 17 00:00:00 2001 From: Michael Lange Date: Thu, 12 Oct 2023 13:23:43 +0000 Subject: [PATCH 3/4] TransformInline: Fix rescoping in expression substitution This would otherwise in-place update associate statements, as they are type definitions. This would create issues for replicated calls when symbols are not rescoped properly. --- loki/transform/transform_inline.py | 4 +- tests/test_transform_inline.py | 62 +++++++++++++++++++++++++++++- 2 files changed, 64 insertions(+), 2 deletions(-) diff --git a/loki/transform/transform_inline.py b/loki/transform/transform_inline.py index ebd30bc9c..91f41b9bb 100644 --- a/loki/transform/transform_inline.py +++ b/loki/transform/transform_inline.py @@ -284,7 +284,9 @@ def _map_unbound_dims(var, val): argmap = recursive_expression_map_update(argmap, max_iterations=10) # Substitute argument calls into a copy of the body - member_body = SubstituteExpressions(argmap).visit(member.body.body) + member_body = SubstituteExpressions(argmap, rebuild_scopes=True).visit( + member.body.body, scope=routine + ) # Inline substituted body within a pair of marker comments comment = Comment(f'! [Loki] inlined member subroutine: {member.name}') diff --git a/tests/test_transform_inline.py b/tests/test_transform_inline.py index 3487bc83a..d1a7b545b 100644 --- a/tests/test_transform_inline.py +++ b/tests/test_transform_inline.py @@ -12,7 +12,7 @@ from conftest import jit_compile, jit_compile_lib, available_frontends from loki import ( Builder, Module, Subroutine, FindNodes, Import, FindVariables, - CallStatement, Loop, BasicType, DerivedType, OMNI + CallStatement, Loop, BasicType, DerivedType, Associate, OMNI ) from loki.ir import Assignment from loki.transform import ( @@ -549,3 +549,63 @@ def test_inline_member_routines_sequence_assoc(frontend): # Expect to fail here due to use of sequence association with pytest.raises(RuntimeError): inline_member_procedures(routine=routine) + + +@pytest.mark.parametrize('frontend', available_frontends()) +def test_inline_member_routines_with_associate(frontend): + """ + Ensure that internal routines with :any:`Associate` constructs get + inlined as expected. + """ + fcode = """ +subroutine acraneb_transt(klon, klev, kidia, kfdia, ktdia) + use parkind1, only: jpim, jprb + implicit none + + integer(kind=jpim), intent(in) :: klon, klev, kidia, kfdia, ktdia + integer(kind=jpim) :: jlon, jlev + + real(kind=jprb) :: zq1(klon) + real(kind=jprb) :: zq2(klon, klev) + + call delta_t(zq1) + + do jlev = ktdia, klev + call delta_t(zq2(1:klon,jlev)) + + enddo + +contains + +subroutine delta_t(pq) + implicit none + + real(kind=jprb), intent(in) :: pq(klon) + real(kind=jprb) :: x, z + + associate(zz => z) + + do jlon = 1,klon + x = x + pq(jlon) + enddo + end associate +end subroutine + +end subroutine acraneb_transt + """ + + routine = Subroutine.from_source(fcode, frontend=frontend) + + inline_member_procedures(routine=routine) + + assert not routine.members + loops = FindNodes(Loop).visit(routine.body) + assert len(loops) == 3 + + assigns = FindNodes(Assignment).visit(routine.body) + assert len(assigns) == 2 + assert assigns[0].rhs == 'x + zq1(jlon)' + assert assigns[1].rhs == 'x + zq2(jlon, jlev)' + + assocs = FindNodes(Associate).visit(routine.body) + assert len(assocs) == 2 From 0de4f6f694f1c84e90bf97d9eea03a856fd70c26 Mon Sep 17 00:00:00 2001 From: Michael Lange Date: Thu, 12 Oct 2023 14:12:43 +0000 Subject: [PATCH 4/4] TransformInline: Fix OMNI test by using more specific kinds --- tests/test_transform_inline.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/tests/test_transform_inline.py b/tests/test_transform_inline.py index d1a7b545b..b5ac37bd0 100644 --- a/tests/test_transform_inline.py +++ b/tests/test_transform_inline.py @@ -559,14 +559,13 @@ def test_inline_member_routines_with_associate(frontend): """ fcode = """ subroutine acraneb_transt(klon, klev, kidia, kfdia, ktdia) - use parkind1, only: jpim, jprb implicit none - integer(kind=jpim), intent(in) :: klon, klev, kidia, kfdia, ktdia - integer(kind=jpim) :: jlon, jlev + integer(kind=4), intent(in) :: klon, klev, kidia, kfdia, ktdia + integer(kind=4) :: jlon, jlev - real(kind=jprb) :: zq1(klon) - real(kind=jprb) :: zq2(klon, klev) + real(kind=8) :: zq1(klon) + real(kind=8) :: zq2(klon, klev) call delta_t(zq1) @@ -580,8 +579,8 @@ def test_inline_member_routines_with_associate(frontend): subroutine delta_t(pq) implicit none - real(kind=jprb), intent(in) :: pq(klon) - real(kind=jprb) :: x, z + real(kind=8), intent(in) :: pq(klon) + real(kind=8) :: x, z associate(zz => z) @@ -601,7 +600,7 @@ def test_inline_member_routines_with_associate(frontend): assert not routine.members loops = FindNodes(Loop).visit(routine.body) assert len(loops) == 3 - + assigns = FindNodes(Assignment).visit(routine.body) assert len(assigns) == 2 assert assigns[0].rhs == 'x + zq1(jlon)'