GlobalVarOffload: globals used to declare array sizes now also offloaded

Author: awnawab

loki/analyse/analyse_dataflow.py

@@ -275,7 +275,9 @@ def visit_Import(self, o, **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)
+        arrays = [a for a in o.symbols if isinstance(a, Array)]
+        uses = set(v for a in arrays for v in FindVariables().visit(a.dimensions))
+        return self.visit_Node(o, defines_symbols=defines, uses_symbols=uses, **kwargs)
 
 
 class DataflowAnalysisDetacher(Transformer):

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')
@@ -428,6 +433,8 @@ def test_global_variable_analysis(frontend, key, config, global_variable_analysi
         if item == 'global_var_analysis_data_mod#some_type':
             continue
         for trafo_data_key, trafo_data_value in item.trafo_data[key].items():
+            print(item)
+            print(trafo_data_value)
             assert (
                 sorted(
                     tuple(str(vv) for vv in v) if isinstance(v, tuple) else str(v)
@@ -465,8 +472,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 +493,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 +502,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 +522,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'}
     }