Merge pull request #88 from ecmwf-ifs/nabr-sgraph

SGraph draft before Scheduler integration

loki/bulk/scheduler.py

@@ -21,7 +21,7 @@
 from loki.module import Module
 
 
-__all__ = ['Scheduler', 'SchedulerConfig']
+__all__ = ['Scheduler', 'SchedulerConfig', 'SGraph']
 
 
 class SchedulerConfig:
@@ -70,7 +70,10 @@ def __init__(self, default, routines, disable=None, dimensions=None, dic2p=None,
 
     @classmethod
     def from_dict(cls, config):
-        default = config['default']
+        """
+        Populate :any:`SchedulerConfig` from the given :any:`dict` :data:`config`
+        """
+        default = config.get('default', {})
         if 'routine' in config:
             config['routines'] = OrderedDict((r['name'], r) for r in config.get('routine', []))
         else:
@@ -98,13 +101,50 @@ def from_dict(cls, config):
 
     @classmethod
     def from_file(cls, path):
+        """
+        Populate :any:`SchedulerConfig` from a toml file at :data:`path`
+        """
         import toml  # pylint: disable=import-outside-toplevel
         # Load configuration file and process options
         with Path(path).open('r') as f:
             config = toml.load(f)
 
         return cls.from_dict(config)
 
+    @staticmethod
+    def match_item_keys(item_name, keys):
+        """
+        Helper routine to match a :any:`Item` name, which includes a scope,
+        to entries in a config property, where names are allowed to appear
+        without the relevant scope names
+        """
+        item_name = item_name.lower()
+        item_names = (item_name, item_name[item_name.find('#')+1:])
+        return tuple(key for key in keys or () if key in item_names)
+
+    def create_item_config(self, name):
+        """
+        Create the bespoke config `dict` for an :any:`Item`
+
+        The resulting config object contains the :attr:`default`
+        values and any item-specific overwrites and additions.
+        """
+        keys = self.match_item_keys(name, self.routines)
+        if len(keys) > 1:
+            if self.default.get('strict'):
+                raise RuntimeError(f'{name} matches multiple config entries: {", ".join(keys)}')
+            warning(f'{name} matches multiple config entries: {", ".join(keys)}')
+        item_conf = self.default.copy()
+        for key in keys:
+            item_conf.update(self.routines[key])
+        return item_conf
+
+    def is_disabled(self, name):
+        """
+        Check if the item with the given :data:`name` is marked as `disabled`
+        """
+        return len(self.match_item_keys(name, self.disable)) > 0
+
 
 class Scheduler:
     """
@@ -788,3 +828,102 @@ def write_cmake_plan(self, filepath, mode, buildpath, rootpath):
 
             s_remove = '\n'.join(f'    {s}' for s in sources_to_remove)
             f.write(f'set( LOKI_SOURCES_TO_REMOVE \n{s_remove}\n   )\n')
+
+
+class SGraph:
+
+    def __init__(self, seed, item_cache, config=None):
+        self._graph = nx.DiGraph()
+        self.populate(seed, item_cache, config)
+
+    def populate(self, seed, item_cache, config):
+        queue = deque()
+
+        # Insert the seed objects
+        for name in as_tuple(seed):
+            if '#' not in name:
+                name = f'#{name}'
+            item = item_cache.get(name)
+
+            if not item:
+                # We may have to create the corresponding module's definitions first
+                module_item = item_cache.get(name[:name.index('#')])
+                if module_item:
+                    module_item.create_definition_items(item_cache=item_cache, config=config)
+                    item = item_cache.get(name)
+
+            if item:
+                self.add_node(item)
+                queue.append(item)
+            else:
+                debug('No item found for seed "%s"', name)
+
+        # Populate the graph
+        while queue:
+            item = queue.popleft()
+
+            if item.expand:
+                dependencies = []
+                items_to_ignore = [*item.block, *item.ignore]
+                for dependency in item.create_dependency_items(item_cache=item_cache, config=config):
+                    if not SchedulerConfig.match_item_keys(dependency.name, items_to_ignore):
+                        dependencies += [dependency]
+                new_items = [item_ for item_ in dependencies if item_ not in self._graph]
+                if new_items:
+                    self.add_nodes(new_items)
+                    queue.extend(new_items)
+                self.add_edges((item, item_) for item_ in dependencies)
+
+    @property
+    def items(self):
+        return tuple(self._graph.nodes)
+
+    @property
+    def dependencies(self):
+        return tuple(self._graph.edges)
+
+    def add_node(self, item):
+        self._graph.add_node(item)
+
+    def add_nodes(self, items):
+        self._graph.add_nodes_from(items)
+
+    def add_edge(self, edge):
+        self._graph.add_edge(edge[0], edge[1])
+
+    def add_edges(self, edges):
+        self._graph.add_edges_from(edges)
+
+    def export_to_file(self, dotfile_path):
+        """
+        Generate a dotfile from the current graph
+
+        Parameters
+        ----------
+        dotfile_path : str or pathlib.Path
+            Path to write the dotfile to. A corresponding graphical representation
+            will be created with an additional ``.pdf`` appendix.
+        """
+        try:
+            import graphviz as gviz  # pylint: disable=import-outside-toplevel
+        except ImportError:
+            warning('[Loki] Failed to load graphviz, skipping file export generation...')
+            return
+
+        path = Path(dotfile_path)
+        graph = gviz.Digraph(format='pdf', strict=True, graph_attr=(('rankdir', 'LR'),))
+
+        # Insert all nodes in the graph
+        style = {
+            'color': 'black', 'shape': 'box', 'fillcolor': 'limegreen', 'style': 'filled'
+        }
+        for item in self.items:
+            graph.node(item.name.upper(), **style)
+
+        # Insert all edges in the schedulers graph
+        graph.edges((a.name.upper(), b.name.upper()) for a, b in self.dependencies)
+
+        try:
+            graph.render(path, view=False)
+        except gviz.ExecutableNotFound as e:
+            warning(f'[Loki] Failed to render callgraph due to graphviz error:\n  {e}')

loki/frontend/regex.py

@@ -42,6 +42,7 @@ class RegexParserClass(Flag):
     pattern matching can be switched on and off for some pattern classes, and thus the overall
     parse time reduced.
     """
+    EmptyClass = 0
     ProgramUnitClass = auto()
     InterfaceClass = auto()
     ImportClass = auto()
@@ -447,7 +448,8 @@ def match(self, reader, parser_classes, scope):
             contains = None
 
         module.__initialize__(  # pylint: disable=unnecessary-dunder-call
-            name=module.name, spec=spec, contains=contains, source=module.source, incomplete=True
+            name=module.name, spec=spec, contains=contains, source=module.source, incomplete=True,
+            parser_classes=parser_classes
         )
 
         if match.span()[0] > 0:
@@ -538,7 +540,8 @@ def match(self, reader, parser_classes, scope):
 
         routine.__initialize__(  # pylint: disable=unnecessary-dunder-call
             name=routine.name, args=routine._dummies, is_function=routine.is_function,
-            prefix=prefix, spec=spec, contains=contains, source=routine.source, incomplete=True
+            prefix=prefix, spec=spec, contains=contains, source=routine.source,
+            incomplete=True, parser_classes=parser_classes
         )
 
         if match.span()[0] > 0:
@@ -897,7 +900,8 @@ class VariableDeclarationPattern(Pattern):
     def __init__(self):
         super().__init__(
             r'^(((?:type|class)[ \t]*\([ \t]*(?P<typename>\w+)[ \t]*\))|' # TYPE or CLASS keyword with typename
-            r'^([ \t]*(?P<basic_type>(logical|real|integer|complex|character))(\((kind|len)=[a-z0-9_-]+\))?[ \t]*))'
+            r'^([ \t]*(?P<basic_type>(logical|real|integer|complex|character))'
+            r'(?P<param>\((kind|len)=[a-z0-9_-]+\))?[ \t]*))'
             r'(?:[ \t]*,[ \t]*[a-z]+(?:\((.(\(.*\))?)*?\))?)*'  # Optional attributes
             r'(?:[ \t]*::)?'  # Optional `::` delimiter
             r'[ \t]*'  # Some white space
@@ -929,6 +933,11 @@ def match(self, reader, parser_classes, scope):
             type_ = SymbolAttributes(BasicType.from_str(match['basic_type']))
         assert type_
 
+        if match['param']:
+            param = match['param'].strip().strip('()').split('=')
+            if len(param) == 1 or param[0].lower() == 'kind':
+                type_ = type_.clone(kind=sym.Variable(name=param[-1], scope=scope))
+
         variables = self._remove_quoted_string_nested_parentheses(match['variables'])  # Remove dimensions
         variables = re.sub(r'=(?:>)?[^,]*(?=,|$)', r'', variables) # Remove initialization
         variables = variables.replace(' ', '').split(',')  # Variable names without white space

loki/module.py

@@ -67,13 +67,15 @@ class Module(ProgramUnit):
         Mark the object as incomplete, i.e. only partially parsed. This is
         typically the case when it was instantiated using the :any:`Frontend.REGEX`
         frontend and a full parse using one of the other frontends is pending.
+    parser_classes : :any:`RegexParserClass`, optional
+        Provide the list of parser classes used during incomplete regex parsing
     """
 
     def __init__(
             self, name=None, docstring=None, spec=None, contains=None,
             default_access_spec=None, public_access_spec=None, private_access_spec=None,
             ast=None, source=None, parent=None, symbol_attrs=None, rescope_symbols=False,
-            incomplete=False
+            incomplete=False, parser_classes=None
     ):
         super().__init__(parent=parent)
 
@@ -84,12 +86,12 @@ def __init__(
             name=name, docstring=docstring, spec=spec, contains=contains,
             default_access_spec=default_access_spec, public_access_spec=public_access_spec,
             private_access_spec=private_access_spec, ast=ast, source=source,
-            rescope_symbols=rescope_symbols, incomplete=incomplete
+            rescope_symbols=rescope_symbols, incomplete=incomplete, parser_classes=parser_classes
         )
 
     def __initialize__(
             self, name=None, docstring=None, spec=None, contains=None,
-            ast=None, source=None, rescope_symbols=False, incomplete=False,
+            ast=None, source=None, rescope_symbols=False, incomplete=False, parser_classes=None,
             default_access_spec=None, public_access_spec=None, private_access_spec=None
     ):
         # Apply dimension pragma annotations to declarations
@@ -110,7 +112,7 @@ def __initialize__(
 
         super().__initialize__(
             name=name, docstring=docstring, spec=spec, contains=contains, ast=ast,
-            source=source, rescope_symbols=rescope_symbols, incomplete=incomplete
+            source=source, rescope_symbols=rescope_symbols, incomplete=incomplete, parser_classes=parser_classes
         )
 
     @classmethod

loki/program_unit.py

@@ -8,7 +8,10 @@
 from abc import abstractmethod
 
 from loki import ir
-from loki.frontend import Frontend, parse_omni_source, parse_ofp_source, parse_fparser_source
+from loki.frontend import (
+    Frontend, parse_omni_source, parse_ofp_source, parse_fparser_source,
+    RegexParserClass
+)
 from loki.scope import Scope
 from loki.tools import CaseInsensitiveDict, as_tuple, flatten
 from loki.types import ProcedureType
@@ -52,16 +55,20 @@ class ProgramUnit(Scope):
         Mark the object as incomplete, i.e. only partially parsed. This is
         typically the case when it was instantiated using the :any:`Frontend.REGEX`
         frontend and a full parse using one of the other frontends is pending.
+    parser_classes : :any:`RegexParserClass`, optional
+        Provide the list of parser classes used during incomplete regex parsing
     """
 
     def __initialize__(self, name, docstring=None, spec=None, contains=None,
-                       ast=None, source=None, rescope_symbols=False, incomplete=False):
+                       ast=None, source=None, rescope_symbols=False, incomplete=False,
+                       parser_classes=None):
         # Common properties
         assert name and isinstance(name, str)
         self.name = name
         self._ast = ast
         self._source = source
         self._incomplete = incomplete
+        self._parser_classes = parser_classes
 
         # Bring arguments into shape
         if spec is not None and not isinstance(spec, ir.Section):
@@ -235,7 +242,11 @@ def make_complete(self, **frontend_args):
         frontend = frontend_args.pop('frontend', Frontend.FP)
         definitions = frontend_args.get('definitions')
         xmods = frontend_args.get('xmods')
-        parser_classes = frontend_args.get('parser_classes')
+        parser_classes = frontend_args.get('parser_classes', RegexParserClass.AllClasses)
+        if frontend == Frontend.REGEX and self._parser_classes:
+            if self._parser_classes == parser_classes:
+                return
+            parser_classes = parser_classes | self._parser_classes
 
         # If this object does not have a parent, we create a temporary parent scope
         # and make sure the node exists in the parent scope. This way, the existing
@@ -462,6 +473,15 @@ def enum_symbols(self):
         """
         return as_tuple(flatten(enum.symbols for enum in FindNodes(ir.Enumeration).visit(self.spec or ())))
 
+    @property
+    def definitions(self):
+        """
+        The list of IR nodes defined by this program unit.
+
+        Returns an empty tuple by default and can be overwritten by derived nodes.
+        """
+        return ()
+
     @property
     def symbols(self):
         """

loki/scope.py

@@ -283,6 +283,18 @@ def rescope_symbols(self):
         from loki.expression import AttachScopes  # pylint: disable=import-outside-toplevel,cyclic-import
         AttachScopes().visit(self)
 
+    def make_complete(self, **frontend_args):
+        """
+        Trigger a re-parse of the object if incomplete to produce a full Loki IR
+
+        See :any:`ProgramUnit.make_complete` for more details.
+
+        This method relays the call only to the :attr:`parent`.
+        """
+        if hasattr(super(), 'make_complete'):
+            super().make_complete(**frontend_args)
+        self.parent.make_complete(**frontend_args)
+
     def clone(self, **kwargs):
         """
         Create a copy of the scope object with the option to override individual