Use generic enrichment in scheduler, scripts and tests

loki/bulk/scheduler.py

@@ -265,6 +265,18 @@ def dependencies(self):
         """
         return as_tuple(self.item_graph.edges)
 
+    @property
+    def definitions(self):
+        """
+        The list of definitions that the source files in the
+        callgraph provide
+        """
+        return tuple(
+            definition
+            for item in self.item_graph
+            for definition in item.source.definitions
+        )
+
     @property
     def file_graph(self):
         """
@@ -541,35 +553,37 @@ def _parse_items(self):
         """
         # Force the parsing of the routines
         build_args = self.build_args.copy()
-        build_args['definitions'] = as_tuple(build_args['definitions'])
+        build_args['definitions'] = as_tuple(build_args['definitions']) + self.definitions
         for item in reversed(list(nx.topological_sort(self.item_graph))):
             item.source.make_complete(**build_args)
-            build_args['definitions'] += item.source.definitions
+
 
     @Timer(logger=perf, text='[Loki::Scheduler] Enriched call tree in {:.2f}s')
     def _enrich(self):
         """
         Enrich subroutine calls for inter-procedural transformations
         """
-        # Force the parsing of the routines in the call tree
+        definitions = self.definitions
         for item in self.item_graph:
             if not isinstance(item, SubroutineItem):
                 continue
 
-            # Enrich with all routines in the call tree
-            item.routine.enrich_calls(routines=self.routines)
-            item.routine.enrich_types(typedefs=self.typedefs)
+            # Enrich all modules and subroutines in the source file with
+            # the definitions of the scheduler's graph
+            for node in item.source.modules + item.source.subroutines:
+                node.enrich(definitions, recurse=True)
 
             # Enrich item with meta-info from outside of the callgraph
-            for routine in item.enrich:
-                lookup_name = self.find_routine(routine)
+            for name in as_tuple(item.enrich):
+                lookup_name = self.find_routine(name)
                 if not lookup_name:
-                    warning(f'Scheduler could not find file for enrichment:\n{routine}')
+                    warning(f'Scheduler could not find file for enrichment:\n{name}')
                     if self.config.default['strict']:
-                        raise FileNotFoundError(f'Source path not found for routine {routine}')
+                        raise FileNotFoundError(f'Source path not found for routine {name}')
                     continue
                 self.obj_map[lookup_name].make_complete(**self.build_args)
-                item.routine.enrich_calls(self.obj_map[lookup_name].all_subroutines)
+                for node in item.source.modules + item.source.subroutines:
+                    node.enrich(self.obj_map[lookup_name].definitions, recurse=True)
 
     def item_successors(self, item):
         """

scripts/loki_transform.py

@@ -332,7 +332,7 @@ def transpile(build, header, source, driver, cpp, include, define, frontend, xmo
                                   frontend=frontend)
     driver = Sourcefile.from_file(driver, xmods=xmod, frontend=frontend)
     # Ensure that the kernel calls have all meta-information
-    driver[driver_name].enrich(routines=kernel[kernel_name])
+    driver[driver_name].enrich(kernel[kernel_name])
 
     kernel_item = SubroutineItem(f'#{kernel_name.lower()}', source=kernel)
     driver_item = SubroutineItem(f'#{driver_name.lower()}', source=driver)

tests/test_scheduler.py

@@ -60,7 +60,8 @@
     fexprgen, Transformation, BasicType, CMakePlanner, Subroutine,
     SubroutineItem, ProcedureBindingItem, gettempdir, ProcedureSymbol,
     ProcedureType, DerivedType, TypeDef, Scalar, Array, FindInlineCalls,
-    Import, Variable, GenericImportItem, GlobalVarImportItem, flatten
+    Import, Variable, GenericImportItem, GlobalVarImportItem, flatten,
+    CaseInsensitiveDict
 )
 
 
@@ -122,6 +123,25 @@ def edges(self):
         return list(self._re_edges.findall(self.text))
 
 
+def test_scheduler_enrichment(here, config, frontend):
+    projA = here/'sources/projA'
+
+    scheduler = Scheduler(
+        paths=projA, includes=projA/'include', config=config,
+        seed_routines=['driverA'], frontend=frontend
+    )
+
+    for item in scheduler.item_graph:
+        if not isinstance(item, SubroutineItem):
+            continue
+        dependency_map = CaseInsensitiveDict(
+            (item_.local_name, item_) for item_ in scheduler.item_successors(item)
+        )
+        for call in FindNodes(CallStatement).visit(item.routine.body):
+            if call_item := dependency_map.get(str(call.name)):
+                assert call.routine is call_item.routine
+
+
 @pytest.mark.skipif(not graphviz_present(), reason='Graphviz is not installed')
 @pytest.mark.parametrize('with_file_graph', [True, False, 'filegraph_simple'])
 def test_scheduler_graph_simple(here, config, frontend, with_file_graph):

transformations/tests/test_single_column.py

@@ -147,7 +147,7 @@ def test_extract_sca_nested_level_zero(frontend, horizontal):
   END SUBROUTINE compute_level_zero
 """, frontend=frontend)
 
-    source['compute_column'].enrich(routines=level_zero.all_subroutines)
+    source['compute_column'].enrich(level_zero.all_subroutines)
 
     # Apply single-column extraction trasnformation in topological order
     sca_transform = ExtractSCATransformation(horizontal=horizontal)
@@ -208,7 +208,7 @@ def test_extract_sca_nested_level_one(frontend, horizontal):
   END SUBROUTINE compute_level_one
 """, frontend=frontend)
 
-    source['compute_column'].enrich(routines=level_one.all_subroutines)
+    source['compute_column'].enrich(level_one.all_subroutines)
 
     # Apply single-column extraction trasnformation in topological order
     sca_transform = ExtractSCATransformation(horizontal=horizontal)