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..b5ac37bd0 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,62 @@ 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) + implicit none + + integer(kind=4), intent(in) :: klon, klev, kidia, kfdia, ktdia + integer(kind=4) :: jlon, jlev + + real(kind=8) :: zq1(klon) + real(kind=8) :: 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=8), intent(in) :: pq(klon) + real(kind=8) :: 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 diff --git a/transformations/tests/test_single_column_coalesced.py b/transformations/tests/test_single_column_coalesced.py index 619a38f52..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( @@ -567,6 +567,16 @@ 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]) + 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) assert driver_loops[6] in FindNodes(Loop).visit(driver_loops[3].body) @@ -596,6 +606,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) 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