Merge pull request #252 from ecmwf-ifs/naml-nabr-the-new-scheduler

Updates for EC-physics regression

loki/bulk/configure.py

@@ -42,6 +42,8 @@ class SchedulerConfig:
         visualisation. These are intended for utility routines that
         pop up in many routines but can be ignored in terms of program
         control flow, like ``flush`` or ``abort``.
+    enable_imports : bool
+        Disable the inclusion of module imports as scheduler dependencies.
     transformation_configs : dict
         Dicts with transformation-specific options
     frontend_args : dict
@@ -50,11 +52,12 @@ class SchedulerConfig:
 
     def __init__(
             self, default, routines, disable=None, dimensions=None,
-            transformation_configs=None, frontend_args=None
+            transformation_configs=None, enable_imports=False, frontend_args=None
     ):
         self.default = default
         self.disable = as_tuple(disable)
         self.dimensions = dimensions
+        self.enable_imports = enable_imports
 
         self.routines = CaseInsensitiveDict(routines)
         self.transformation_configs = transformation_configs
@@ -74,8 +77,7 @@ def from_dict(cls, config):
         default = config.get('default', {})
         routines = config.get('routines', [])
         disable = default.get('disable', None)
-        if 'imports' in default:
-            warning('The enable_imports option in the scheduler config is deprecated. Imports are always enabled.')
+        enable_imports = default.get('enable_imports', False)
 
         # Add any dimension definitions contained in the config dict
         dimensions = config.get('dimensions', {})
@@ -91,7 +93,8 @@ def from_dict(cls, config):
 
         return cls(
             default=default, routines=routines, disable=disable, dimensions=dimensions,
-            transformation_configs=transformation_configs, frontend_args=frontend_args
+            transformation_configs=transformation_configs, frontend_args=frontend_args,
+            enable_imports=enable_imports
         )
 
     @classmethod

loki/bulk/scheduler.py

@@ -201,12 +201,15 @@ def _discover(self):
             }
             self.item_factory.item_cache.update(definition_items)
 
+        # (Re-)build the SGraph after discovery for later traversals
+        self._sgraph = SGraph.from_seed(self.seeds, self.item_factory, self.config)
+
     @property
     def sgraph(self):
         """
         Create and return the :any:`SGraph` constructed from the :attr:`seeds` of the Scheduler.
         """
-        return SGraph.from_seed(self.seeds, self.item_factory, self.config)
+        return self._sgraph
 
     @property
     def items(self):
@@ -244,7 +247,10 @@ def file_graph(self):
         :any:`SGraph`
             A dependency graph containing only :any:`FileItem` nodes
         """
-        return self.sgraph.as_filegraph(self.item_factory, self.config)
+        item_filter = None if self.config.enable_imports else ProcedureItem
+        return self.sgraph.as_filegraph(
+            self.item_factory, self.config, item_filter=item_filter
+        )
 
     def __getitem__(self, name):
         """
@@ -267,10 +273,13 @@ def _parse_items(self):
         # Force the parsing of the routines
         default_frontend_args = self.build_args.copy()
         default_frontend_args['definitions'] = as_tuple(default_frontend_args['definitions']) + self.definitions
-        for item in SFilter(self.sgraph.as_filegraph(self.item_factory, self.config), reverse=True):
+        for item in SFilter(self.file_graph, reverse=True):
             frontend_args = self.config.create_frontend_args(item.name, default_frontend_args)
             item.source.make_complete(**frontend_args)
 
+        # Re-build the SGraph after parsing to pick up all new connections
+        self._sgraph = SGraph.from_seed(self.seeds, self.item_factory, self.config)
+
     @Timer(logger=info, text='[Loki::Scheduler] Enriched call tree in {:.2f}s')
     def _enrich(self):
         """
@@ -448,6 +457,7 @@ def _get_definition_items(_item, sgraph_items):
 
         if transformation.creates_items:
             self._discover()
+
             self._parse_items()
 
     def callgraph(self, path, with_file_graph=False, with_legend=False):
@@ -576,7 +586,9 @@ def write_cmake_plan(self, filepath, mode, buildpath, rootpath):
         sources_to_remove = []
         sources_to_transform = []
 
-        for item in self.items:
+        # Filter the SGraph to get a pure call-tree
+        item_filter = None if self.config.enable_imports else ProcedureItem
+        for item in SFilter(self.sgraph, item_filter=item_filter):
             if item.is_ignored:
                 continue
 
@@ -602,7 +614,6 @@ def write_cmake_plan(self, filepath, mode, buildpath, rootpath):
                     sources_to_append += [newsource]
                     sources_to_remove += [sourcepath]
 
-        info(f'[Loki] CMakePlanner writing plan: {filepath}')
         with Path(filepath).open('w') as f:
             s_transform = '\n'.join(f'    {s}' for s in sources_to_transform)
             f.write(f'set( LOKI_SOURCES_TO_TRANSFORM \n{s_transform}\n   )\n')
@@ -631,6 +642,7 @@ def __init__(self):
         self._graph = nx.DiGraph()
 
     @classmethod
+    @Timer(logger=info, text='[Loki::Scheduler] Built SGraph from seed in {:.2f}s')
     def from_seed(cls, seed, item_factory, config=None):
         """
         Create a new :any:`SGraph` using :data:`seed` as starting point.
@@ -748,7 +760,8 @@ def _add_children(self, item, item_factory, config, dependencies=None):
         if new_items:
             self.add_nodes(new_items)
 
-        self.add_edges((item, item_) for item_ in dependencies)
+        # Careful not to include cycles (from recursive TypeDefs)
+        self.add_edges((item, item_) for item_ in dependencies if not item == item_)
         return new_items
 
     def _populate(self, seed, item_factory, config):

loki/sourcefile.py

@@ -22,7 +22,7 @@
 
 )
 from loki.ir import Section, RawSource, Comment, PreprocessorDirective
-from loki.logging import info, perf
+from loki.logging import info, debug, perf
 from loki.module import Module
 from loki.program_unit import ProgramUnit
 from loki.subroutine import Subroutine
@@ -362,11 +362,12 @@ def make_complete(self, **frontend_args):
         if not self._incomplete:
             return
 
+        frontend = frontend_args.pop('frontend', FP)
+
         log = f'[Loki::Sourcefile] Finished constructing from {self.path}' + ' in {:.2f}s'
-        with Timer(logger=info, text=log):
+        with Timer(logger=debug if frontend == REGEX else info, text=log):
 
             # Sanitize frontend_args
-            frontend = frontend_args.pop('frontend', FP)
             if isinstance(frontend, str):
                 frontend = Frontend[frontend.upper()]
             if frontend == REGEX:

tests/test_cmake.py

@@ -63,7 +63,7 @@ def fixture_config(here):
     the file path
     """
     default_config = {
-        'default': {'role': 'kernel', 'expand': True, 'strict': True},
+        'default': {'role': 'kernel', 'expand': True, 'strict': True, 'enable_imports': True},
         'routines': {
             'driverB': {'role': 'driver'},
         }

tests/test_lint/test_linter.py

@@ -478,7 +478,8 @@ def output(self, handler_reports):
                 'role': 'kernel',
                 'expand': True,
                 'strict': False,
-                'block': ['compute_l2']
+                'block': ['compute_l2'],
+                'enable_imports': True,
             },
             'routines': routines
         }

tests/test_scheduler.py

@@ -86,7 +86,8 @@ def fixture_config():
             'role': 'kernel',
             'expand': True,
             'strict': True,
-            'disable': ['abort']
+            'disable': ['abort'],
+            'enable_imports': True,
         },
         'routines': {}
     }
@@ -995,15 +996,19 @@ def test_scheduler_dependencies_ignore(here, preprocess, frontend):
     projB = here/'sources/projB'
 
     configA = SchedulerConfig.from_dict({
-        'default': {'role': 'kernel', 'expand': True, 'strict': True},
+        'default': {
+            'role': 'kernel', 'expand': True, 'strict': True, 'enable_imports': True
+        },
         'routines': {
             'driverB': {'role': 'driver'},
             'kernelB': {'ignore': ['ext_driver']},
         }
     })
 
     configB = SchedulerConfig.from_dict({
-        'default': {'role': 'kernel', 'expand': True, 'strict': True},
+        'default': {
+            'role': 'kernel', 'expand': True, 'strict': True, 'enable_imports': True
+        },
         'routines': {
             'ext_driver': {'role': 'kernel'}
         }
@@ -1387,7 +1392,7 @@ def test_scheduler_typebound(here, config, frontend, proj_typebound_dependencies
 
     scheduler = Scheduler(
         paths=proj, seed_routines=['driver'], config=config,
-        full_parse=False, frontend=frontend
+        full_parse=False, frontend=frontend,
     )
 
     assert set(scheduler.items) == set(proj_typebound_dependencies)

transformations/tests/test_data_offload.py

@@ -37,6 +37,7 @@ def fixture_config():
             'role': 'kernel',
             'expand': True,
             'strict': True,
+            'enable_imports': True,
         },
     }
 

transformations/tests/test_pool_allocator.py

@@ -173,6 +173,7 @@ def test_pool_allocator_temporaries(frontend, generate_driver_stack, block_dim,
             'role': 'kernel',
             'expand': True,
             'strict': True,
+            'enable_imports': True,
         },
         'routines': {
             'driver': {'role': 'driver'}
@@ -505,7 +506,8 @@ def test_pool_allocator_temporaries_kernel_sequence(frontend, block_dim, directi
             'mode': 'idem',
             'role': 'kernel',
             'expand': True,
-            'strict': True
+            'strict': True,
+            'enable_imports': True,
         },
         'routines': {
             'driver': {'role': 'driver'}
@@ -789,7 +791,8 @@ def test_pool_allocator_temporaries_kernel_nested(frontend, block_dim, directive
             'mode': 'idem',
             'role': 'kernel',
             'expand': True,
-            'strict': True
+            'strict': True,
+            'enable_imports': True,
         },
         'routines': {
             'driver': {'role': 'driver', 'real_kind': 'jwrb'}
@@ -1020,7 +1023,8 @@ def test_pool_allocator_more_call_checks(frontend, block_dim, caplog):
             'mode': 'idem',
             'role': 'kernel',
             'expand': True,
-            'strict': True
+            'strict': True,
+            'enable_imports': True,
         },
         'routines': {
             'kernel': {}
@@ -1151,7 +1155,8 @@ def test_pool_allocator_args_vs_kwargs(frontend, block_dim):
             'role': 'kernel',
             'expand': True,
             'strict': True,
-            'disable': ['parkind1']
+            'disable': ['parkind1'],
+            'enable_imports': True,
         },
         'routines': {
             'driver': {'role': 'driver'}

transformations/tests/test_transform_derived_types.py

@@ -33,6 +33,7 @@ def fixture_config():
             'role': 'kernel',
             'expand': True,
             'strict': True,
+            'enable_imports': True,
         },
         'routines': {
             'driver': {