extend hoist variables functionality

loki/transformations/hoist_variables.py

@@ -88,7 +88,8 @@
     SubstituteExpressions, is_dimension_constant
 )
 from loki.ir import (
-    CallStatement, Allocation, Deallocation, Transformer, FindNodes, Comment, Import
+    CallStatement, Allocation, Deallocation, Transformer, FindNodes, Comment, Import,
+    Assignment
 )
 from loki.tools.util import is_iterable, as_tuple, CaseInsensitiveDict, flatten
 
@@ -204,12 +205,17 @@ class HoistVariablesTransformation(Transformation):
     ----------
     as_kwarguments : boolean
         Whether to pass the hoisted arguments as `args` or `kwargs`.
+    remap_dimensions : boolean
+        Remap dimensions based on variables that are used for initializing
+        other variables that could end up as dimensions for hoisted arrays.
+        Thus, account for possibly uninitialized variables used as dimensions.
     """
 
     _key = 'HoistVariablesTransformation'
 
-    def __init__(self, as_kwarguments=False):
+    def __init__(self, as_kwarguments=False, remap_dimensions=True):
         self.as_kwarguments = as_kwarguments
+        self.remap_dimensions = remap_dimensions
 
     def transform_subroutine(self, routine, **kwargs):
         """
@@ -242,7 +248,12 @@ def transform_subroutine(self, routine, **kwargs):
                                f'the correct key.')
 
         if role == 'driver':
-            self.driver_variable_declaration(routine, item.trafo_data[self._key]["to_hoist"])
+            if self.remap_dimensions:
+                to_hoist = self.driver_variable_declaration_dim_remapping(routine,
+                        item.trafo_data[self._key]["to_hoist"])
+            else:
+                to_hoist = item.trafo_data[self._key]["to_hoist"]
+            self.driver_variable_declaration(routine, to_hoist)
         else:
             # We build the list of temporaries that are hoisted to the calling routine
             # Because this requires adding an intent, we need to make sure they are not
@@ -319,6 +330,33 @@ def driver_variable_declaration(self, routine, variables):
         """
         routine.variables += tuple(v.rescope(routine) for v in variables)
 
+    @staticmethod
+    def driver_variable_declaration_dim_remapping(routine, variables):
+        """
+        Take a list of variables and remap their dimensions for those being
+        arrays to account for possibly uninitialized variables/dimensions. 
+
+        Parameters
+        ----------
+        routine : :any:`Subroutine`
+            The relevant subroutine.
+        variables : tuple of :any:`Variable`
+            The tuple of variables for remapping.
+        """
+        dim_vars = [
+            dim_var
+            for var in variables if isinstance(var, sym.Array)
+            for dim_var in FindVariables().visit(var.dimensions)
+        ]
+        dim_map = {
+            assignment.lhs: assignment.rhs
+            for assignment in FindNodes(Assignment).visit(routine.body)
+            if assignment.lhs in dim_vars
+        }
+        variables = [var.clone(dimensions=SubstituteExpressions(dim_map).visit(var.dimensions))
+                if isinstance(var, sym.Array) else var for var in variables]
+        return variables
+
     def driver_call_argument_remapping(self, routine, call, variables):
         """
         Callback method to re-map hoisted arguments for the driver-level routine.

loki/transformations/tests/test_hoist_variables.py

@@ -672,3 +672,72 @@ def test_hoist_mixed_variable_declarations(tmp_path, frontend, config):
     imports = FindNodes(ir.Import).visit(scheduler['#driver'].ir.spec)
     assert len(imports) == 2
     assert 'n' in scheduler['#driver'].ir.imported_symbols
+
+
+@pytest.mark.parametrize('frontend', available_frontends())
+@pytest.mark.parametrize('remap_dimensions', (False, True))
+def test_hoist_dim_mapping(tmp_path, frontend, config, remap_dimensions):
+
+    fcode_driver = """
+subroutine driver(NLON, NB, FIELD1)
+    use kernel_mod, only: kernel
+    implicit none
+    INTEGER, INTENT(IN) :: NLON, NB
+    integer :: b
+    integer, intent(inout) :: field1(nlon, nb)
+    integer :: local_nlon
+    local_nlon = nlon
+    do b=1,nb
+        call KERNEL(local_nlon, field1(:,b))
+    end do
+end subroutine driver
+    """.strip()
+    fcode_kernel = """
+module kernel_mod
+    implicit none
+contains
+    subroutine kernel(klon, field1)
+        implicit none
+        integer, intent(in) :: klon
+        integer, intent(inout) :: field1(klon)
+        integer :: tmp1(klon)
+        integer :: jl
+
+        do jl=1,klon
+            tmp1(jl) = 0
+            field1(jl) = tmp1(jl)
+        end do
+
+    end subroutine kernel
+end module kernel_mod
+    """.strip()
+
+    (tmp_path/'driver.F90').write_text(fcode_driver)
+    (tmp_path/'kernel_mod.F90').write_text(fcode_kernel)
+
+    config = {
+        'default': {
+            'mode': 'idem',
+            'role': 'kernel',
+            'expand': True,
+            'strict': True,
+            'enable_imports': True
+        },
+        'routines': {
+            'driver': {'role': 'driver'}
+        }
+    }
+
+    scheduler = Scheduler(
+        paths=[tmp_path], config=SchedulerConfig.from_dict(config), frontend=frontend, xmods=[tmp_path]
+    )
+
+    scheduler.process(transformation=HoistTemporaryArraysAnalysis())
+    scheduler.process(transformation=HoistVariablesTransformation(remap_dimensions=remap_dimensions))
+
+    driver_var_map = scheduler['#driver'].ir.variable_map
+    assert 'kernel_tmp1' in driver_var_map
+    if remap_dimensions:
+        assert driver_var_map['kernel_tmp1'].dimensions == ('nlon',)
+    else:
+        assert driver_var_map['kernel_tmp1'].dimensions == ('local_nlon',)