Started on: making Loki plan trafo (pipeline) aware

loki/batch/item.py

@@ -517,6 +517,12 @@ def path(self):
         """
         return self.source.path
 
+    @property
+    def orig_path(self):
+        """
+        The filepath of the associated source file
+        """
+        return self.source.orig_path
 
 class FileItem(Item):
     """

loki/batch/scheduler.py

@@ -508,6 +508,120 @@ def _get_definition_items(_item, sgraph_items):
             self._discover()
             self._parse_items()
 
+    def process_plan(self, transformation):
+        """
+        """
+        if isinstance(transformation, Transformation):
+            self.process_plan_transformation(transformation=transformation)
+
+        elif isinstance(transformation, Pipeline):
+            self.process_plan_pipeline(pipeline=transformation)
+
+        else:
+            error('[Loki::Scheduler] Batch processing requires Transformation or Pipeline object')
+            raise RuntimeError('[Loki] Could not batch process {transformation_or_pipeline}')
+
+    def process_plan_pipeline(self, pipeline):
+        """
+        Process a given :any:`Pipeline` by applying its assocaited
+        transformations in turn.
+
+        Parameters
+        ----------
+        transformation : :any:`Pipeline`
+            The transformation pipeline to apply
+        """
+        for transformation in pipeline.transformations:
+            self.process_plan_transformation(transformation)
+
+    def process_plan_transformation(self, transformation):
+        """
+        Process all :attr:`items` in the scheduler's graph
+
+        By default, the traversal is performed in topological order, which
+        ensures that an item is processed before the items it depends upon
+        (e.g., via a procedure call)
+        This order can be reversed in the :any:`Transformation` manifest by
+        setting :any:`Transformation.reverse_traversal` to ``True``.
+
+        The scheduler applies the transformation to the scope corresponding to
+        each item in the scheduler's graph, determined by the :any:`Item.scope_ir`
+        property. For example, for a :any:`ProcedureItem`, the transformation is
+        applied to the corresponding :any:`Subroutine` object.
+
+        Optionally, the traversal can be performed on a source file level only,
+        if the transformation has set :any:`Transformation.traverse_file_graph`
+        to ``True``. This uses the :attr:`filegraph` to process the dependency tree.
+        If combined with a :any:`Transformation.item_filter`, only source files with
+        at least one object corresponding to an item of that type are processed.
+
+        Parameters
+        ----------
+        transformation : :any:`Transformation`
+            The transformation to apply over the dependency tree
+        """
+        def _get_definition_items(_item, sgraph_items):
+            # For backward-compatibility with the DependencyTransform and LinterTransformation
+            if not transformation.traverse_file_graph:
+                return None
+
+            # Recursively obtain all definition items but exclude any that are not part of the original SGraph
+            items = ()
+            for item in _item.create_definition_items(item_factory=self.item_factory, config=self.config):
+                # Recursion gives us only items that are included in the SGraph, or the parent scopes
+                # of items included in the SGraph
+                child_items = _get_definition_items(item, sgraph_items)
+                # If the current item has relevant children, or is included in the SGraph itself, we
+                # include it in the list of items
+                if child_items or item in sgraph_items:
+                    if transformation.process_ignored_items or not item.is_ignored:
+                        items += (item,) + child_items
+            return items
+
+        trafo_name = transformation.__class__.__name__
+        log = f'[Loki::Scheduler] Applied transformation <{trafo_name}>' + ' in {:.2f}s'
+        with Timer(logger=info, text=log):
+
+            # Extract the graph iteration properties from the transformation
+            item_filter = as_tuple(transformation.item_filter)
+            if transformation.traverse_file_graph:
+                sgraph = self.sgraph
+                graph = sgraph.as_filegraph(
+                    self.item_factory, self.config, item_filter=item_filter,
+                    exclude_ignored=not transformation.process_ignored_items
+                )
+                sgraph_items = sgraph.items
+                traversal = SFilter(
+                    graph, reverse=transformation.reverse_traversal,
+                    include_external=self.config.default.get('strict', True)
+                )
+            else:
+                graph = self.sgraph
+                sgraph_items = graph.items
+                traversal = SFilter(
+                    graph, item_filter=item_filter, reverse=transformation.reverse_traversal,
+                    exclude_ignored=not transformation.process_ignored_items,
+                    include_external=self.config.default.get('strict', True)
+                )
+
+            for _item in traversal:
+                if isinstance(_item, ExternalItem):
+                    raise RuntimeError(f'Cannot apply {trafo_name} to {_item.name}: Item is marked as external.')
+
+                transformation.apply_plan(
+                    _item.scope_ir, role=_item.role, mode=_item.mode,
+                    item=_item, targets=_item.targets, items=_get_definition_items(_item, sgraph_items),
+                    successors=graph.successors(_item, item_filter=item_filter),
+                    depths=graph.depths, build_args=self.build_args # , item_factory=self.item_factory
+                )
+
+        if transformation.renames_items:
+            self.rekey_item_cache()
+
+        if transformation.creates_items:
+            self._discover()
+            self._parse_items()
+
     def callgraph(self, path, with_file_graph=False, with_legend=False):
         """
         Generate a callgraph visualization and dump to file.
@@ -640,16 +754,18 @@ def write_cmake_plan(self, filepath, mode, buildpath, rootpath):
             if item.is_ignored:
                 continue
 
-            sourcepath = item.path.resolve()
-            newsource = sourcepath.with_suffix(f'.{mode.lower()}.F90')
+            sourcepath = item.orig_path.resolve()
+            newsource = item.path.resolve()
+            # sourcepath = item.path.resolve()
+            # newsource = sourcepath.with_suffix(f'.{mode.lower()}.F90')
             if buildpath:
                 newsource = buildpath/newsource.name
-
-            # Make new CMake paths relative to source again
+            # 
+            # # Make new CMake paths relative to source again
             if rootpath is not None:
                 sourcepath = sourcepath.relative_to(rootpath)
-
-            debug(f'Planning:: {item.name} (role={item.role}, mode={mode})')
+            # 
+            # debug(f'Planning:: {item.name} (role={item.role}, mode={mode})')
 
             # Inject new object into the final binary libs
             if newsource not in sources_to_append:

loki/batch/transformation.py

@@ -412,3 +412,122 @@ def post_apply_module(self, module, rescope_symbols):
         # Ensure all objects in the IR are in the module's scope.
         if rescope_symbols:
             module.rescope_symbols()
+
+    def plan_subroutine(self, source, **kwargs):
+        """
+        ...
+        """
+
+    def plan_module(self, source, **kwargs):
+        """
+        ...
+        """
+
+    def plan_file(self, source, **kwargs):
+        """
+        ...
+        """
+
+    def apply_plan(self, source, **kwargs):
+        """
+        ...
+        """
+        if isinstance(source, Sourcefile):
+            self.apply_plan_file(source, **kwargs)
+
+        if isinstance(source, Subroutine):
+            self.apply_plan_subroutine(source, **kwargs)
+
+        if isinstance(source, Module):
+            self.apply_plan_module(source, **kwargs)
+
+    def apply_plan_file(self, sourcefile, **kwargs):
+        """
+        """
+        if not isinstance(sourcefile, Sourcefile):
+            raise TypeError('Transformation.apply_file can only be applied to Sourcefile object')
+
+        # if sourcefile._incomplete:
+        #     raise RuntimeError('Transformation.apply_file requires Sourcefile to be complete')
+
+        item = kwargs.pop('item', None)
+        items = kwargs.pop('items', None)
+        role = kwargs.pop('role', None)
+        targets = kwargs.pop('targets', None)
+
+        # Apply file-level transformations
+        self.plan_file(sourcefile, item=item, role=role, targets=targets, items=items, **kwargs)
+
+        # Recurse to modules, if configured
+        if self.recurse_to_modules:
+            if items:
+                # Recursion into all module items in the current file
+                for item in items:
+                    if isinstance(item, ModuleItem):
+                        # Currently, we don't get the role for modules correct as 'driver'
+                        # if the role overwrite in the config marks only specific procedures
+                        # as driver, but everything else as kernel by default. This is in particular the
+                        # case, if the ModuleWrapTransformation is applied to a driver routine.
+                        # For that reason, we set the role as unspecified (None) if not the role is
+                        # universally equal throughout the module
+                        item_role = item.role
+                        definitions_roles = {_it.role for _it in items if _it.scope_name == item.name}
+                        if definitions_roles != {item_role}:
+                            item_role = None
+
+                        # Provide the list of items that belong to this module
+                        item_items = tuple(_it for _it in items if _it.scope is item.ir)
+
+                        self.plan_module(
+                            item.ir, item=item, role=item_role, targets=item.targets, items=item_items, **kwargs
+                        )
+            else:
+                for module in sourcefile.modules:
+                    self.plan_module(module, item=item, role=role, targets=targets, items=items, **kwargs)
+
+        # Recurse into procedures, if configured
+        if self.recurse_to_procedures:
+            if items:
+                # Recursion into all subroutine items in the current file
+                for item in items:
+                    if isinstance(item, ProcedureItem):
+                        self.plan_subroutine(
+                            item.ir, item=item, role=item.role, targets=item.targets, **kwargs
+                        )
+            else:
+                for routine in sourcefile.all_subroutines:
+                    self.plan_subroutine(routine, item=item, role=role, targets=targets, **kwargs)
+
+    def apply_plan_subroutine(self, subroutine, **kwargs):
+        """
+        """
+        if not isinstance(subroutine, Subroutine):
+            raise TypeError('Transformation.apply_subroutine can only be applied to Subroutine object')
+
+        # if subroutine._incomplete:
+        #     raise RuntimeError('Transformation.apply_subroutine requires Subroutine to be complete')
+
+        # Apply the actual transformation for subroutines
+        self.plan_subroutine(subroutine, **kwargs)
+
+        # Recurse to internal procedures
+        if self.recurse_to_internal_procedures:
+            for routine in subroutine.subroutines:
+                self.apply_plan_subroutine(routine, **kwargs)
+
+    def apply_plan_module(self, module, **kwargs):
+        """
+        """
+        if not isinstance(module, Module):
+            raise TypeError('Transformation.apply_module can only be applied to Module object')
+
+        # if module._incomplete:
+        #     raise RuntimeError('Transformation.apply_module requires Module to be complete')
+
+        # Apply the actual transformation for modules
+        self.transform_module(module, **kwargs)
+
+        # Recurse to procedures contained in this module
+        if self.recurse_to_procedures:
+            for routine in module.subroutines:
+                self.apply_plan_subroutine(routine, **kwargs)

loki/sourcefile.py

@@ -64,6 +64,7 @@ class Sourcefile:
 
     def __init__(self, path, ir=None, ast=None, source=None, incomplete=False, parser_classes=None):
         self.path = Path(path) if path is not None else path
+        self.orig_path = self.path
         if ir is not None and not isinstance(ir, Section):
             ir = Section(body=ir)
         self.ir = ir

loki/transformations/build_system/file_write.py

@@ -73,3 +73,21 @@ def transform_file(self, sourcefile, **kwargs):
         if build_args and (output_dir := build_args.get('output_dir', None)) is not None:
             sourcepath = Path(output_dir)/sourcepath.name
         sourcefile.write(path=sourcepath, cuf=self.cuf)
+
+    def plan_file(self, sourcefile, **kwargs):
+        item = kwargs.get('item', None)
+        if not item and 'items' in kwargs:
+            if kwargs['items']:
+                item = kwargs['items'][0]
+
+        if not item:
+            raise ValueError('No Item provided; required to determine file write path')
+
+        _mode = item.mode if item.mode else 'loki'
+        _mode = _mode.replace('-', '_')  # Sanitize mode string
+
+        path = Path(item.path)
+        suffix = self.suffix if self.suffix else path.suffix
+        sourcepath = Path(item.path).with_suffix(f'.{_mode}{suffix}')
+        item.source.path = sourcepath
+

loki/transformations/parametrise.py

@@ -82,6 +82,8 @@
     scheduler.process(transformation=ParametriseTransformation(dic2p=dic2p, replace_by_value=True))
 """
 
+from pathlib import Path
+from loki import ProcedureItem
 from loki.batch import Transformation
 from loki.expression import symbols as sym
 from loki.ir import nodes as ir, Transformer, FindNodes
@@ -90,9 +92,74 @@
 from loki.transformations.inline import inline_constant_parameters
 
 
-__all__ = ['ParametriseTransformation']
+__all__ = ['ParametriseTransformation', 'DuplicateKernel', 'RemoveKernel']
 
 
+class DuplicateKernel(Transformation):
+
+    def __init__(self, kernels=None):
+        self.kernels = tuple(kernel.lower() for kernel in as_tuple(kernels))
+
+    def transform_subroutine(self, routine, **kwargs):
+        calls = FindNodes(ir.CallStatement).visit(routine.body)
+        call_map = {}
+        for call in calls:
+            if str(call.name).lower() in self.kernels:
+                call_map[call] = (call, call.clone())
+        routine.body = Transformer(call_map).visit(routine.body)
+
+    def plan_subroutine(self, routine, **kwargs):
+        item = kwargs.get('item', None)
+        if not item and 'items' in kwargs:
+            if kwargs['items']:
+                item = kwargs['items'][0]
+
+        successors = as_tuple(kwargs.get('successors'))
+
+        # this only renames, however we want to create a duplicated kernel and not only as duplicate call ...
+        """
+        for child in successors:
+            if not isinstance(child, ProcedureItem):
+                continue
+            print(f"plan_subroutine - child for {routine} : {child}")
+            if child.local_name.lower() in self.kernels:
+                path = Path(child.path)
+                suffix = path.suffix
+                child.source.path = Path(child.path).with_suffix(f'.duplicate{suffix}') 
+        """
+
+class RemoveKernel(Transformation):
+
+    def __init__(self, kernels=None):
+        self.kernels = tuple(kernel.lower() for kernel in as_tuple(kernels))
+
+    def transform_subroutine(self, routine, **kwargs):
+        calls = FindNodes(ir.CallStatement).visit(routine.body)
+        call_map = {}
+        for call in calls:
+            if str(call.name).lower() in self.kernels:
+                call_map[call] = None
+        routine.body = Transformer(call_map).visit(routine.body)
+
+    """
+    def plan_subroutine(self, sourcefile, **kwargs):
+        item = kwargs.get('item', None)
+        if not item and 'items' in kwargs:
+            if kwargs['items']:
+                item = kwargs['items'][0]
+
+        if not item:
+            raise ValueError('No Item provided; required to determine file write path')
+
+        _mode = item.mode if item.mode else 'loki'
+        _mode = _mode.replace('-', '_')  # Sanitize mode string
+
+        path = Path(item.path)
+        suffix = self.suffix if self.suffix else path.suffix
+        sourcepath = Path(item.path).with_suffix(f'.{_mode}{suffix}')
+        item.source.path = sourcepath
+    """
+
 class ParametriseTransformation(Transformation):
     """
     Parametrise variables with provided values.