Merge pull request #320 from ecmwf-ifs/naan-dataflow-fixes

`DataFlowAnalysis` bug fixes

Author: reuterbal

loki/analyse/analyse_dataflow.py

@@ -10,7 +10,7 @@
 """
 
 from contextlib import contextmanager
-from loki.expression import FindVariables, Array, FindInlineCalls
+from loki.expression import FindVariables, Array, FindInlineCalls, ProcedureSymbol, FindTypedSymbols
 from loki.tools import as_tuple, flatten
 from loki.types import BasicType
 from loki.ir import Visitor, Transformer
@@ -232,9 +232,9 @@ def visit_CallStatement(self, o, **kwargs):
             outvals = [val for arg, val in o.arg_iter() if str(arg.type.intent).lower() in ('inout', 'out')]
             invals = [val for arg, val in o.arg_iter() if str(arg.type.intent).lower() in ('inout', 'in')]
 
+            arrays = [v for v in FindVariables().visit(outvals) if isinstance(v, Array)]
+            dims = set(v for a in arrays for v in self._symbols_from_expr(a.dimensions))
             for val in outvals:
-                arrays = [v for v in FindVariables().visit(outvals) if isinstance(v, Array)]
-                dims = set(v for a in arrays for v in FindVariables().visit(a.dimensions))
                 exprs = self._symbols_from_expr(val)
                 defines |= {e for e in exprs if not e in dims}
                 uses |= dims
@@ -269,12 +269,14 @@ def visit_Deallocation(self, o, **kwargs):
     visit_Nullify = visit_Deallocation
 
     def visit_Import(self, o, **kwargs):
-        defines = self._symbols_from_expr(o.symbols or ())
+        defines = set(s.clone(dimensions=None) for s in FindTypedSymbols().visit(o.symbols or ())
+                      if isinstance(s, ProcedureSymbol))
         return self.visit_Node(o, defines_symbols=defines, **kwargs)
 
     def visit_VariableDeclaration(self, o, **kwargs):
         defines = self._symbols_from_expr(o.symbols, condition=lambda v: v.type.initial is not None)
-        return self.visit_Node(o, defines_symbols=defines, **kwargs)
+        uses = {v for a in o.symbols if isinstance(a, Array) for v in self._symbols_from_expr(a.dimensions)}
+        return self.visit_Node(o, defines_symbols=defines, uses_symbols=uses, **kwargs)
 
 
 class DataflowAnalysisDetacher(Transformer):

loki/analyse/tests/test_analyse_dataflow.py

@@ -10,7 +10,7 @@
 from loki import (
     Subroutine, FindNodes, Assignment, Loop, Conditional, Pragma, fgen, Sourcefile,
     CallStatement, MultiConditional, MaskedStatement, ProcedureSymbol, WhileLoop,
-    Associate
+    Associate, Module
 )
 from loki.analyse import (
     dataflow_analysis_attached, read_after_write_vars, loop_carried_dependencies
@@ -281,6 +281,41 @@ def test_analyse_interface(frontend):
         assert isinstance(list(routine.spec.defines_symbols)[0], ProcedureSymbol)
         assert 'random_call' in routine.spec.defines_symbols
 
+
+@pytest.mark.parametrize('frontend', available_frontends())
+def test_analyse_imports(frontend):
+    fcode_module = """
+module some_mod
+implicit none
+real :: my_global
+contains
+subroutine random_call(v_out,v_in,v_inout)
+
+  real,intent(in)  :: v_in
+  real,intent(out)  :: v_out
+  real,intent(inout)  :: v_inout
+
+
+end subroutine random_call
+end module some_mod
+""".strip()
+
+    fcode = """
+subroutine test()
+use some_mod, only: my_global, random_call
+implicit none
+
+end subroutine test
+""".strip()
+
+    module = Module.from_source(fcode_module, frontend=frontend)
+    routine = Subroutine.from_source(fcode, frontend=frontend, definitions=module)
+
+    with dataflow_analysis_attached(routine):
+        assert len(routine.spec.defines_symbols) == 1
+        assert 'random_call' in routine.spec.defines_symbols
+
+
 @pytest.mark.parametrize('frontend', available_frontends())
 def test_analyse_enriched_call(frontend):
     fcode = """
@@ -429,26 +464,30 @@ def test_analyse_call_args_array_slicing(frontend):
 
 end subroutine random_call
 
-subroutine test(v,n)
+subroutine test(v,n,b)
 implicit none
 
   integer,intent(out) :: v(:)
   integer,intent( in) :: n
+  integer,intent( in) :: b(n)
 
   call random_call(v(n))
+  call random_call(v(b(1)))
 
 end subroutine test
     """.strip()
 
     source = Sourcefile.from_source(fcode, frontend=frontend)
     routine = source['test']
 
-    call = FindNodes(CallStatement).visit(routine.body)[0]
+    calls = FindNodes(CallStatement).visit(routine.body)
     routine.enrich(source.all_subroutines)
 
     with dataflow_analysis_attached(routine):
-        assert 'n' in call.uses_symbols
-        assert not 'n' in call.defines_symbols
+        assert 'n' in calls[0].uses_symbols
+        assert not 'n' in calls[0].defines_symbols
+        assert 'b' in calls[1].uses_symbols
+        assert not 'b' in calls[0].defines_symbols
 
 
 @pytest.mark.parametrize('frontend', available_frontends())

loki/transformations/data_offload.py

@@ -296,6 +296,10 @@ def transform_subroutine(self, routine, **kwargs):
                 var for var in routine.body.uses_symbols
                 if var.name in import_map or (var.parent and var.parents[0].name in import_map)
             }
+            uses_imported_symbols |= {
+                var for var in routine.spec.uses_symbols
+                if var.name in import_map or (var.parent and var.parents[0].name in import_map)
+            }
             defines_imported_symbols = {
                 var for var in routine.body.defines_symbols
                 if var.name in import_map or (var.parent and var.parents[0].name in import_map)

loki/transformations/tests/test_data_offload.py

@@ -256,6 +256,8 @@ def fixture_global_variable_analysis_code(tmp_path):
     integer, parameter :: nval = 5
     integer, parameter :: nfld = 3
 
+    integer :: n
+
     integer :: iarr(nfld)
     real :: rarr(nval, nfld)
 end module global_var_analysis_header_mod
@@ -297,10 +299,11 @@ def fixture_global_variable_analysis_code(tmp_path):
 
 contains
     subroutine kernel_a(arg, tt)
-        use global_var_analysis_header_mod, only: iarr, nval, nfld
+        use global_var_analysis_header_mod, only: iarr, nval, nfld, n
 
         real, intent(inout) :: arg(:,:)
         type(some_type), intent(in) :: tt
+        real :: tmp(n)
         integer :: i, j
 
         do i=1,nfld
@@ -390,7 +393,7 @@ def test_global_variable_analysis(frontend, key, config, global_variable_analysi
 
     expected_trafo_data = {
         'global_var_analysis_header_mod': {
-            'declares': {f'iarr({nfld_dim})', f'rarr({nval_dim}, {nfld_dim})'},
+            'declares': {f'iarr({nfld_dim})', f'rarr({nval_dim}, {nfld_dim})', 'n'},
             'offload': {}
         },
         'global_var_analysis_data_mod': {
@@ -402,6 +405,7 @@ def test_global_variable_analysis(frontend, key, config, global_variable_analysi
             'defines_symbols': set(),
             'uses_symbols': nval_data | nfld_data | {
                 (f'iarr({nfld_dim})', 'global_var_analysis_header_mod'),
+                ('n', 'global_var_analysis_header_mod'),
                 (f'rarr({nval_dim}, {nfld_dim})', 'global_var_analysis_header_mod')
             }
         },
@@ -416,6 +420,7 @@ def test_global_variable_analysis(frontend, key, config, global_variable_analysi
             'defines_symbols': {('rdata(:, :, :)', 'global_var_analysis_data_mod')},
             'uses_symbols': nval_data | nfld_data | {
                 ('rdata(:, :, :)', 'global_var_analysis_data_mod'),
+                ('n', 'global_var_analysis_header_mod'),
                 ('tt', 'global_var_analysis_data_mod'), ('tt%vals', 'global_var_analysis_data_mod'),
                 (f'iarr({nfld_dim})', 'global_var_analysis_header_mod'),
                 (f'rarr({nval_dim}, {nfld_dim})', 'global_var_analysis_header_mod')
@@ -465,8 +470,8 @@ def test_global_variable_offload(frontend, key, config, global_variable_analysis
 
     expected_trafo_data = {
         'global_var_analysis_header_mod': {
-            'declares': {f'iarr({nfld_dim})', f'rarr({nval_dim}, {nfld_dim})'},
-            'offload': {f'iarr({nfld_dim})', f'rarr({nval_dim}, {nfld_dim})'}
+            'declares': {f'iarr({nfld_dim})', f'rarr({nval_dim}, {nfld_dim})', 'n'},
+            'offload': {f'iarr({nfld_dim})', f'rarr({nval_dim}, {nfld_dim})', 'n'}
         },
         'global_var_analysis_data_mod': {
             'declares': {'rdata(:, :, :)', 'tt'},
@@ -486,7 +491,7 @@ def test_global_variable_offload(frontend, key, config, global_variable_analysis
 
     # Verify imports have been added to the driver
     expected_imports = {
-        'global_var_analysis_header_mod': {'iarr', 'rarr'},
+        'global_var_analysis_header_mod': {'iarr', 'rarr', 'n'},
         'global_var_analysis_data_mod': {'rdata'}
     }
 
@@ -495,7 +500,7 @@ def test_global_variable_offload(frontend, key, config, global_variable_analysis
         assert {var.name.lower() for var in import_.symbols} == expected_imports[import_.module.lower()]
 
     expected_h2d_pragmas = {
-        'update device': {'iarr', 'rdata', 'rarr'},
+        'update device': {'iarr', 'rdata', 'rarr', 'n'},
         'enter data copyin': {'tt%vals'}
     }
     expected_d2h_pragmas = {
@@ -515,7 +520,7 @@ def test_global_variable_offload(frontend, key, config, global_variable_analysis
 
     # Verify declarations have been added to the header modules
     expected_declarations = {
-        'global_var_analysis_header_mod': {'iarr', 'rarr'},
+        'global_var_analysis_header_mod': {'iarr', 'rarr', 'n'},
         'global_var_analysis_data_mod': {'rdata', 'tt'}
     }