Merge pull request #321 from ecmwf-ifs/naan-import-var-sizes

Add transformation generated imports

Author: reuterbal

loki/transformations/hoist_variables.py

@@ -80,15 +80,17 @@
         scheduler.process(transformation=HoistTemporaryArraysTransformationAllocatable(key=key))
 """
 
+from collections import defaultdict
+
 from loki.batch import Transformation, ProcedureItem
 from loki.expression import (
     symbols as sym, FindVariables, FindInlineCalls,
     SubstituteExpressions, is_dimension_constant
 )
 from loki.ir import (
-    CallStatement, Allocation, Deallocation, Transformer, FindNodes
+    CallStatement, Allocation, Deallocation, Transformer, FindNodes, Comment, Import
 )
-from loki.tools.util import is_iterable, as_tuple, CaseInsensitiveDict
+from loki.tools.util import is_iterable, as_tuple, CaseInsensitiveDict, flatten
 
 from loki.transformations.utilities import single_variable_declaration
 
@@ -139,9 +141,14 @@ def transform_subroutine(self, routine, **kwargs):
         if role != 'driver':
             variables = self.find_variables(routine)
             item.trafo_data[self._key]["to_hoist"] = variables
+            dims = flatten([getattr(v, 'shape', []) for v in variables])
+            import_map = routine.import_map
+            item.trafo_data[self._key]["imported_sizes"] = [(d.type.module, d) for d in dims
+                                                            if str(d) in import_map]
             item.trafo_data[self._key]["hoist_variables"] = [var.clone(name=f'{routine.name}_{var.name}')
                                                              for var in variables]
         else:
+            item.trafo_data[self._key]["imported_sizes"] = []
             item.trafo_data[self._key]["to_hoist"] = []
             item.trafo_data[self._key]["hoist_variables"] = []
 
@@ -166,6 +173,7 @@ def transform_subroutine(self, routine, **kwargs):
             item.trafo_data[self._key]["hoist_variables"].extend(hoist_variables)
             item.trafo_data[self._key]["hoist_variables"] = list(dict.fromkeys(
                 item.trafo_data[self._key]["hoist_variables"]))
+            item.trafo_data[self._key]["imported_sizes"] += child.trafo_data[self._key]["imported_sizes"]
 
     def find_variables(self, routine):
         """
@@ -273,6 +281,25 @@ def transform_subroutine(self, routine, **kwargs):
                         routine=routine, call=call, variables=hoisted_variables
                     )
 
+        # Add imports used to define hoisted
+        missing_imports_map = defaultdict(set)
+        import_map = routine.import_map
+        for module, var in item.trafo_data[self._key]["imported_sizes"]:
+            if not var.name in import_map:
+                missing_imports_map[module] |= {var}
+
+        if missing_imports_map:
+            routine.spec.prepend(Comment(text=(
+                '![Loki::HoistVariablesTransformation] ---------------------------------------'
+            )))
+            for module, variables in missing_imports_map.items():
+                routine.spec.prepend(Import(module=module.name, symbols=variables))
+
+            routine.spec.prepend(Comment(text=(
+                '![Loki::HoistVariablesTransformation] '
+                '-------- Added hoisted temporary size imports -------------------------------'
+            )))
+
         routine.body = Transformer(call_map).visit(routine.body)
 
     def driver_variable_declaration(self, routine, variables):

loki/transformations/inline.py

@@ -15,7 +15,7 @@
 from loki.batch import Transformation
 from loki.ir import (
     Import, Comment, Assignment, VariableDeclaration, CallStatement,
-    Transformer, FindNodes, pragmas_attached, is_loki_pragma
+    Transformer, FindNodes, pragmas_attached, is_loki_pragma, Interface
 )
 from loki.expression import (
     symbols as sym, FindVariables, FindInlineCalls, FindLiterals,
@@ -199,9 +199,9 @@ def map_inline_call(self, expr, *args, **kwargs):
 
 def resolve_sequence_association_for_inlined_calls(routine, inline_internals, inline_marked):
     """
-    Resolve sequence association in calls to all member procedures (if `inline_internals = True`)
-    or in calls to procedures that have been marked with an inline pragma (if `inline_marked = True`).
-    If both `inline_internals` and `inline_marked` are `False`, no processing is done.
+    Resolve sequence association in calls to all member procedures (if ``inline_internals = True``)
+    or in calls to procedures that have been marked with an inline pragma (if ``inline_marked = True``).
+    If both ``inline_internals`` and ``inline_marked`` are ``False``, no processing is done.
     """
     call_map = {}
     with pragmas_attached(routine, node_type=CallStatement):
@@ -635,13 +635,26 @@ def inline_marked_subroutines(routine, allowed_aliases=None, adjust_imports=True
 
                 # If we're importing the same module, check for missing symbols
                 if m := imported_module_map.get(impt.module):
-                    if not all(s in m.symbols for s in impt.symbols):
+                    _m = import_map.get(m, m)
+                    if not all(s in _m.symbols for s in impt.symbols):
                         new_symbols = tuple(s.rescope(routine) for s in impt.symbols)
-                        import_map[m] = m.clone(symbols=tuple(set(m.symbols + new_symbols)))
+                        import_map[m] = m.clone(symbols=tuple(set(_m.symbols + new_symbols)))
 
         # Finally, apply the import remapping
         routine.spec = Transformer(import_map).visit(routine.spec)
 
+        # Add missing explicit interfaces from inlined subroutines
+        new_intfs = []
+        intf_symbols = routine.interface_symbols
+        for callee in call_sets.keys():
+            for intf in callee.interfaces:
+                for s in intf.symbols:
+                    if not s in intf_symbols:
+                        new_intfs += [s.type.dtype.procedure,]
+
+        if new_intfs:
+            routine.spec.append(Interface(body=as_tuple(new_intfs)))
+
         # Add Fortran imports to the top, and C-style interface headers at the bottom
         c_imports = tuple(im for im in new_imports if im.c_import)
         f_imports = tuple(im for im in new_imports if not im.c_import)

loki/transformations/pool_allocator.py

@@ -731,10 +731,14 @@ def apply_pool_allocator_to_temporaries(self, routine, item=None):
                     stack_size, stack_storage)
             allocations += allocation
 
-            # Store type information of temporary allocation
-            if item and (_kind := arr.type.kind):
-                if _kind in routine.imported_symbols:
-                    item.trafo_data[self._key]['kind_imports'][_kind] = routine.import_map[_kind.name].module.lower()
+            # Store type and size information of temporary allocation
+            if item:
+                if (kind := arr.type.kind):
+                    if kind in routine.imported_symbols:
+                        item.trafo_data[self._key]['kind_imports'][kind] = routine.import_map[kind.name].module.lower()
+                dims = [d for d in arr.shape if d in routine.imported_symbols]
+                for d in dims:
+                    item.trafo_data[self._key]['kind_imports'][d] = routine.import_map[d.name].module.lower()
 
         routine.spec.append(declarations)
         routine.body.prepend(allocations)

loki/transformations/tests/test_hoist_variables.py

@@ -557,6 +557,11 @@ def test_hoist_mixed_variable_declarations(tmp_path, frontend, config):
     subroutine kernel(start, end, klon, klev, nclv, field1, field2)
         use, intrinsic :: iso_c_binding, only : c_size_t
         implicit none
+        interface
+           subroutine another_kernel(klev)
+               integer, intent(in) :: klev
+           end subroutine another_kernel
+        end interface
         integer, parameter :: jprb = selected_real_kind(13,300)
         integer, intent(in) :: nclv
         integer, intent(in) :: start, end, klon, klev
@@ -582,19 +587,38 @@ def test_hoist_mixed_variable_declarations(tmp_path, frontend, config):
              tmp5(jl, jm, :) = field1(jl)
            enddo
         enddo
+
+        call another_kernel(klev)
     end subroutine kernel
 end module kernel_mod
     """.strip()
+    fcode_mod = """
+module size_mod
+   implicit none
+   integer :: n
+end module size_mod
+""".strip()
+    fcode_another_kernel = """
+subroutine another_kernel(klev)
+    use size_mod, only : n
+    implicit none
+    integer, intent(in) :: klev
+    real :: another_tmp(klev,n)
+end subroutine another_kernel
+""".strip()
 
     (tmp_path/'driver.F90').write_text(fcode_driver)
     (tmp_path/'kernel_mod.F90').write_text(fcode_kernel)
+    (tmp_path/'size_mod.F90').write_text(fcode_mod)
+    (tmp_path/'another_kernel.F90').write_text(fcode_another_kernel)
 
     config = {
         'default': {
             'mode': 'idem',
             'role': 'kernel',
             'expand': True,
-            'strict': True
+            'strict': True,
+            'enable_imports': True
         },
         'routines': {
             'driver': {'role': 'driver'}
@@ -613,11 +637,12 @@ def test_hoist_mixed_variable_declarations(tmp_path, frontend, config):
     driver_variables = (
         'jprb', 'nlon', 'nz', 'nb', 'b',
         'field1(nlon, nb)', 'field2(nlon, nz, nb)',
-        'kernel_tmp2(:,:)', 'kernel_tmp5(:,:,:)'
+        'kernel_tmp2(:,:)', 'kernel_tmp5(:,:,:)', 'another_kernel_another_tmp(:,:)'
     )
     kernel_arguments = (
        'start', 'end', 'klon', 'klev', 'nclv',
-        'field1(klon)', 'field2(klon,klev)', 'tmp2(klon,klev)', 'tmp5(klon,nclv,klev)'
+        'field1(klon)', 'field2(klon,klev)', 'tmp2(klon,klev)', 'tmp5(klon,nclv,klev)',
+        'another_kernel_another_tmp(klev,n)'
     )
 
     # Check hoisting and declaration in driver
@@ -629,8 +654,17 @@ def test_hoist_mixed_variable_declarations(tmp_path, frontend, config):
     assert len(calls) == 1
     assert calls[0].arguments == (
         '1', 'nlon', 'nlon', 'nz', '2', 'field1(:,b)', 'field2(:,:,b)',
-        'kernel_tmp2', 'kernel_tmp5'
+        'kernel_tmp2', 'kernel_tmp5', 'another_kernel_another_tmp'
     )
 
     # Check that fgen works
     assert scheduler['kernel_mod#kernel'].source.to_fortran()
+
+    # Check that imports were updated
+    imports = FindNodes(ir.Import).visit(scheduler['kernel_mod#kernel'].ir.spec)
+    assert len(imports) == 2
+    assert 'n' in scheduler['kernel_mod#kernel'].ir.imported_symbols
+
+    imports = FindNodes(ir.Import).visit(scheduler['#driver'].ir.spec)
+    assert len(imports) == 2
+    assert 'n' in scheduler['#driver'].ir.imported_symbols

loki/transformations/tests/test_inline.py

@@ -738,11 +738,19 @@ def test_inline_marked_subroutines(frontend, adjust_imports):
 
 contains
   subroutine add_one(a)
+    interface
+      subroutine do_something()
+      end subroutine do_something
+    end interface
     real(kind=8), intent(inout) :: a
     a = a + 1
   end subroutine add_one
 
   subroutine add_a_to_b(a, b, n)
+    interface
+      subroutine do_something_else()
+      end subroutine do_something_else
+    end interface
     real(kind=8), intent(inout) :: a(:), b(:)
     integer, intent(in) :: n
     integer :: i
@@ -786,6 +794,14 @@ def test_inline_marked_subroutines(frontend, adjust_imports):
     else:
         assert imports[0].symbols == ('add_one', 'add_a_to_b')
 
+    if adjust_imports:
+        # check that explicit interfaces were imported
+        intfs = driver.interfaces
+        assert len(intfs) == 1
+        assert all(isinstance(s, sym.ProcedureSymbol) for s in driver.interface_symbols)
+        assert 'do_something' in driver.interface_symbols
+        assert 'do_something_else' in driver.interface_symbols
+
 
 @pytest.mark.parametrize('frontend', available_frontends())
 def test_inline_marked_subroutines_with_interfaces(frontend):
@@ -1361,6 +1377,7 @@ def test_inline_transformation_adjust_imports(frontend):
 module bnds_module
   integer :: m
   integer :: n
+  integer :: l
 end module bnds_module
     """
 
@@ -1374,10 +1391,13 @@ def test_inline_transformation_adjust_imports(frontend):
 subroutine test_inline_outer(a, b)
   use bnds_module, only: n
   use test_inline_mod, only: test_inline_inner
+  use test_inline_another_mod, only: test_inline_another_inner
   implicit none
 
   real(kind=8), intent(inout) :: a(n), b(n)
 
+  !$loki inline
+  call test_inline_another_inner()
   !$loki inline
   call test_inline_inner(a, b)
 end subroutine test_inline_outer
@@ -1403,11 +1423,25 @@ def test_inline_transformation_adjust_imports(frontend):
 end subroutine test_inline_inner
 end module test_inline_mod
     """
+
+    fcode_another_inner = """
+module test_inline_another_mod
+  implicit none
+  contains
+
+subroutine test_inline_another_inner()
+  use BNDS_module, only: n, m, l
+
+end subroutine test_inline_another_inner
+end module test_inline_another_mod
+    """
+
     _ = Module.from_source(fcode_another, frontend=frontend)
     _ = Module.from_source(fcode_module, frontend=frontend)
     inner = Module.from_source(fcode_inner, frontend=frontend)
+    another_inner = Module.from_source(fcode_another_inner, frontend=frontend)
     outer = Subroutine.from_source(
-        fcode_outer, definitions=inner, frontend=frontend
+        fcode_outer, definitions=(inner, another_inner), frontend=frontend
     )
 
     trafo = InlineTransformation(
@@ -1430,7 +1464,7 @@ def test_inline_transformation_adjust_imports(frontend):
     assert imports[0].module == 'another_module'
     assert imports[0].symbols == ('x',)
     assert imports[1].module == 'bnds_module'
-    assert 'm' in imports[1].symbols and 'n' in imports[1].symbols
+    assert all(_ in imports[1].symbols for _ in ['l', 'm', 'n'])
 
 
 @pytest.mark.parametrize('frontend', available_frontends())

loki/transformations/tests/test_pool_allocator.py

@@ -1196,6 +1196,8 @@ def test_pool_allocator_args_vs_kwargs(tmp_path, frontend, block_dim_alt, cray_p
     type geom_type
        type(dim_type) :: blk_dim
     end type geom_type
+
+    integer :: n
 end module geom_mod
 """
 
@@ -1228,6 +1230,7 @@ def test_pool_allocator_args_vs_kwargs(tmp_path, frontend, block_dim_alt, cray_p
 contains
     subroutine kernel(start, end, klon, klev, field1, field2)
         use parkind1, only : jpim, jplm
+        use geom_mod, only : n
         implicit none
         integer, parameter :: jwrb = selected_real_kind(13,300)
         integer, intent(in) :: start, end, klon, klev
@@ -1237,6 +1240,7 @@ def test_pool_allocator_args_vs_kwargs(tmp_path, frontend, block_dim_alt, cray_p
         real(kind=jwrb) :: tmp2(klon, klev)
         integer(kind=jpim) :: tmp3(klon*2)
         logical(kind=jplm) :: tmp4(klev)
+        logical(kind=jplm) :: tmp5(klev,n)
         integer :: jk, jl
 
         do jk=1,klev
@@ -1247,6 +1251,7 @@ def test_pool_allocator_args_vs_kwargs(tmp_path, frontend, block_dim_alt, cray_p
             end do
             field1(jl) = tmp1(jl)
             tmp4(jk) = .true.
+            tmp5(jk,1:n) = .true.
         end do
 
         do jl=start,end
@@ -1344,3 +1349,6 @@ def test_pool_allocator_args_vs_kwargs(tmp_path, frontend, block_dim_alt, cray_p
     # check stack size allocation
     allocations = FindNodes(Allocation).visit(driver.body)
     assert len(allocations) == 1 and 'zstack(istsz,geom%blk_dim%nb)' in allocations[0].variables
+
+    # check that array size was imported to the driver
+    assert 'n' in driver.imported_symbols