Merge branch 'main' of github.com:ecmwf-ifs/loki into sbrm-fix-scalar…

…-syntax-for-array-arguments

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}')

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

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)
 
 

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