CMake Plan: Pipeline dry-run and transformation-based plan creation

loki/batch/scheduler.py

@@ -21,7 +21,7 @@
 
 from loki.frontend import FP, REGEX, RegexParserClass
 from loki.tools import as_tuple, CaseInsensitiveDict, flatten
 from loki.logging import info, perf, warning, debug, error
 
 
 __all__ = ['Scheduler']
@@ -382,7 +382,7 @@
             if item.name not in deleted_keys
         )
 
-    def process(self, transformation):
+    def process(self, transformation, plan=False):
         """
         Process all :attr:`items` in the scheduler's graph with either
         a :any:`Pipeline` or a single :any:`Transformation`.
@@ -398,16 +398,16 @@
             The transformation or transformation pipeline to apply
         """
         if isinstance(transformation, Transformation):
-            self.process_transformation(transformation=transformation)
+            self.process_transformation(transformation=transformation, plan=plan)
 
         elif isinstance(transformation, Pipeline):
-            self.process_pipeline(pipeline=transformation)
+            self.process_pipeline(pipeline=transformation, plan=plan)
 
         else:
             error('[Loki::Scheduler] Batch processing requires Transformation or Pipeline object')
             raise RuntimeError('[Loki] Could not batch process {transformation_or_pipeline}')
 
-    def process_pipeline(self, pipeline):
+    def process_pipeline(self, pipeline, plan=False):
         """
         Process a given :any:`Pipeline` by applying its assocaited
         transformations in turn.
@@ -418,9 +418,9 @@
             The transformation pipeline to apply
         """
         for transformation in pipeline.transformations:
-            self.process_transformation(transformation)
+            self.process_transformation(transformation, plan=plan)
 
-    def process_transformation(self, transformation):
+    def process_transformation(self, transformation, plan=False):
         """
         Process all :attr:`items` in the scheduler's graph
 
@@ -498,7 +498,8 @@
                     _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
+                    depths=graph.depths, build_args=self.build_args,
+                    plan=plan
                 )
 
         if transformation.renames_items:
@@ -609,7 +610,7 @@
                 warning(f'[Loki] Failed to render filegraph due to graphviz error:\n  {e}')
 
     @Timer(logger=perf, text='[Loki::Scheduler] Wrote CMake plan file in {:.2f}s')
-    def write_cmake_plan(self, filepath, mode, buildpath, rootpath):
+    def write_cmake_plan(self, filepath, rootpath=None):
         """
         Generate the "plan file" for CMake
 
@@ -628,53 +629,7 @@
         """
         info(f'[Loki] Scheduler writing CMake plan: {filepath}')
 
-        rootpath = None if rootpath is None else Path(rootpath).resolve()
-        buildpath = None if buildpath is None else Path(buildpath)
-        sources_to_append = []
-        sources_to_remove = []
-        sources_to_transform = []
-
-        # Filter the SGraph to get a pure call-tree
-        item_filter = ProcedureItem
-        if self.config.enable_imports:
-            item_filter = as_tuple(item_filter) + (ModuleItem,)
-        graph = self.sgraph.as_filegraph(
-            self.item_factory, self.config, item_filter=item_filter,
-            exclude_ignored=True
-        )
-        traversal = SFilter(graph, reverse=False, include_external=False)
-        for item in traversal:
-            if item.is_ignored:
-                continue
-
-            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
-            if rootpath is not None:
-                sourcepath = sourcepath.relative_to(rootpath)
-
-            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:
-                sources_to_transform += [sourcepath]
-                if item.replicate:
-                    # Add new source file next to the old one
-                    sources_to_append += [newsource]
-                else:
-                    # Replace old source file to avoid ghosting
-                    sources_to_append += [newsource]
-                    sources_to_remove += [sourcepath]
-
-        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')
-
-            s_append = '\n'.join(f'    {s}' for s in sources_to_append)
-            f.write(f'set( LOKI_SOURCES_TO_APPEND \n{s_append}\n   )\n')
-
-            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')
+        from loki.transformations.build_system.plan import CMakePlanTransformation
+        planner = CMakePlanTransformation(rootpath=rootpath)
+        self.process(planner, plan=True)
+        planner.write_plan(filepath)

loki/batch/tests/test_scheduler.py

@@ -74,7 +74,7 @@
     nodes as ir, FindNodes, FindInlineCalls, FindVariables
 )
 from loki.transformations import (
-    DependencyTransformation, ModuleWrapTransformation
+    DependencyTransformation, ModuleWrapTransformation, FileWriteTransformation
 )
 
 
@@ -1001,7 +1001,7 @@
     # Check processing with missing items
     class CheckApply(Transformation):
 
         def apply(self, source, post_apply_rescope_symbols=False, **kwargs):
             assert 'item' in kwargs
             assert not isinstance(kwargs['item'], ExternalItem)
             super().apply(source, post_apply_rescope_symbols=post_apply_rescope_symbols, **kwargs)
@@ -1156,35 +1156,35 @@
     proj_b = sourcedir/'projB'
 
     config = SchedulerConfig.from_dict({
-        'default': {'role': 'kernel', 'expand': True, 'strict': True, 'ignore': ('header_mod',)},
+        'default': {
+            'role': 'kernel',
+            'expand': True,
+            'strict': True,
+            'ignore': ('header_mod',),
+            'mode': 'foobar'
+        },
         'routines': {
             'driverB': {'role': 'driver'},
             'kernelB': {'ignore': ['ext_driver']},
         }
     })
+    builddir = tmp_path/'scheduler_cmake_planner_dummy_dir'
+    builddir.mkdir(exist_ok=True)
 
     # Populate the scheduler
     # (this is the same as SchedulerA in test_scheduler_dependencies_ignore, so no need to
     # check scheduler set-up itself)
     scheduler = Scheduler(
         paths=[proj_a, proj_b], includes=proj_a/'include',
-        config=config, frontend=frontend, xmods=[tmp_path]
+        config=config, frontend=frontend, xmods=[tmp_path],
+        output_dir=builddir
     )
 
     # Apply the transformation
-    builddir = tmp_path/'scheduler_cmake_planner_dummy_dir'
-    builddir.mkdir(exist_ok=True)
     planfile = builddir/'loki_plan.cmake'
 
-    scheduler.write_cmake_plan(
-        filepath=planfile, mode='foobar', buildpath=builddir, rootpath=sourcedir
-    )
-
-    # Validate the generated lists
-    expected_files = {
-        proj_a/'module/driverB_mod.f90', proj_a/'module/kernelB_mod.F90',
-        proj_a/'module/compute_l1_mod.f90', proj_a/'module/compute_l2_mod.f90'
-    }
+    scheduler.process(FileWriteTransformation(), plan=True)
+    scheduler.write_cmake_plan(filepath=planfile, rootpath=sourcedir)
 
     # Validate the plan file content
     plan_pattern = re.compile(r'set\(\s*(\w+)\s*(.*?)\s*\)', re.DOTALL)

loki/batch/transformation.py

@@ -128,6 +128,11 @@ def transform_subroutine(self, routine, **kwargs):
             Keyword arguments for the transformation.
         """
 
+    def plan_subroutine(self, routine, **kwargs):
+        """
+        ...
+        """
+
     def transform_module(self, module, **kwargs):
         """
         Defines the transformation to apply to :any:`Module` items.
@@ -144,6 +149,11 @@ def transform_module(self, module, **kwargs):
             Keyword arguments for the transformation.
         """
 
+    def plan_module(self, module, **kwargs):
+        """
+        ...
+        """
+
     def transform_file(self, sourcefile, **kwargs):
         """
         Defines the transformation to apply to :any:`Sourcefile` items.
@@ -160,7 +170,12 @@ def transform_file(self, sourcefile, **kwargs):
             Keyword arguments for the transformation.
         """
 
-    def apply(self, source, post_apply_rescope_symbols=False, **kwargs):
+    def plan_file(self, sourcefile, **kwargs):
+        """
+        ...
+        """
+
+    def apply(self, source, post_apply_rescope_symbols=False, plan=False, **kwargs):
         """
         Dispatch method to apply transformation to :data:`source`.
 
@@ -179,17 +194,18 @@ def apply(self, source, post_apply_rescope_symbols=False, **kwargs):
             actual transformation.
         """
         if isinstance(source, Sourcefile):
-            self.apply_file(source, **kwargs)
+            self.apply_file(source, plan=plan, **kwargs)
 
         if isinstance(source, Subroutine):
-            self.apply_subroutine(source, **kwargs)
+            self.apply_subroutine(source, plan=plan, **kwargs)
 
         if isinstance(source, Module):
-            self.apply_module(source, **kwargs)
+            self.apply_module(source, plan=plan, **kwargs)
 
-        self.post_apply(source, rescope_symbols=post_apply_rescope_symbols)
+        if not plan:
+            self.post_apply(source, rescope_symbols=post_apply_rescope_symbols)
 
-    def apply_file(self, sourcefile, **kwargs):
+    def apply_file(self, sourcefile, plan=False, **kwargs):
         """
         Apply transformation to all items in :data:`sourcefile`.
 
@@ -212,16 +228,19 @@ def apply_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.transform_file(sourcefile, item=item, role=role, targets=targets, items=items, **kwargs)
+        if plan:
+            self.plan_file(sourcefile, item=item, role=role, targets=targets, items=items, **kwargs)
+        else:
+            if not plan and sourcefile._incomplete:
+                raise RuntimeError('Transformation.apply_file requires Sourcefile to be complete')
+
+            self.transform_file(sourcefile, item=item, role=role, targets=targets, items=items, **kwargs)
 
         # Recurse to modules, if configured
         if self.recurse_to_modules:
@@ -243,27 +262,43 @@ def apply_file(self, sourcefile, **kwargs):
                         # 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.transform_module(
-                            item.ir, item=item, role=item_role, targets=item.targets, items=item_items, **kwargs
-                        )
+                        if plan:
+                            self.plan_module(
+                                item.ir, item=item, role=item_role, targets=item.targets, items=item_items, **kwargs
+                            )
+                        else:
+                            self.transform_module(
+                                item.ir, item=item, role=item_role, targets=item.targets, items=item_items, **kwargs
+                            )
             else:
                 for module in sourcefile.modules:
-                    self.transform_module(module, item=item, role=role, targets=targets, items=items, **kwargs)
+                    if plan:
+                        self.plan_module(module, item=item, role=role, targets=targets, items=items, **kwargs)
+                    else:
+                        self.transform_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.transform_subroutine(
-                            item.ir, item=item, role=item.role, targets=item.targets, **kwargs
-                        )
+                        if plan:
+                            self.plan_subroutine(
+                                item.ir, item=item, role=item.role, targets=item.targets, **kwargs
+                            )
+                        else:
+                            self.transform_subroutine(
+                                item.ir, item=item, role=item.role, targets=item.targets, **kwargs
+                            )
             else:
                 for routine in sourcefile.all_subroutines:
-                    self.transform_subroutine(routine, item=item, role=role, targets=targets, **kwargs)
+                    if plan:
+                        self.plan_subroutine(routine, item=item, role=role, targets=targets, **kwargs)
+                    else:
+                        self.transform_subroutine(routine, item=item, role=role, targets=targets, **kwargs)
 
-    def apply_subroutine(self, subroutine, **kwargs):
+    def apply_subroutine(self, subroutine, plan=False, **kwargs):
         """
         Apply transformation to a given :any:`Subroutine` object and its members.
 
@@ -284,18 +319,21 @@ def apply_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.transform_subroutine(subroutine, **kwargs)
+        if plan:
+            self.plan_subroutine(subroutine, **kwargs)
+        else:
+            if subroutine._incomplete:
+                raise RuntimeError('Transformation.apply_subroutine requires Subroutine to be complete')
+
+            self.transform_subroutine(subroutine, **kwargs)
 
         # Recurse to internal procedures
         if self.recurse_to_internal_procedures:
             for routine in subroutine.subroutines:
-                self.apply_subroutine(routine, **kwargs)
+                self.apply_subroutine(routine, plan=plan, **kwargs)
 
-    def apply_module(self, module, **kwargs):
+    def apply_module(self, module, plan=False, **kwargs):
         """
         Apply transformation to a given :any:`Module` object and its members.
 
@@ -315,16 +353,19 @@ def apply_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)
+        if plan:
+            self.plan_module(module, **kwargs)
+        else:
+            if module._incomplete:
+                raise RuntimeError('Transformation.apply_module requires Module to be complete')
+
+            self.transform_module(module, **kwargs)
 
         # Recurse to procedures contained in this module
         if self.recurse_to_procedures:
             for routine in module.subroutines:
-                self.apply_subroutine(routine, **kwargs)
+                self.apply_subroutine(routine, plan=plan, **kwargs)
 
     def post_apply(self, source, rescope_symbols=False):
         """