Skip to content

Commit

Permalink
TransformInline: Fix rescoping in expression substitution
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
mlange05 committed Oct 12, 2023
1 parent ede5135 commit 186ff67
Show file tree
Hide file tree
Showing 2 changed files with 64 additions and 2 deletions.
4 changes: 3 additions & 1 deletion loki/transform/transform_inline.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}')
Expand Down
62 changes: 61 additions & 1 deletion tests/test_transform_inline.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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

Check failure on line 604 in tests/test_transform_inline.py

View workflow job for this annotation

GitHub Actions / code checks (3.11)

C0303: Trailing whitespace (trailing-whitespace)
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

0 comments on commit 186ff67

Please sign in to comment.