WIP: Plan mode for transformations

loki/batch/scheduler.py

@@ -382,7 +382,7 @@ def rekey_item_cache(self):
             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 @@ def process(self, transformation):
             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 @@ def process_pipeline(self, pipeline):
             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 @@ def _get_definition_items(_item, sgraph_items):
                     _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:
@@ -628,53 +629,7 @@ def write_cmake_plan(self, filepath, mode, buildpath, rootpath):
         """
         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/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):
         """

loki/transformations/build_system/file_write.py

@@ -54,12 +54,7 @@ def item_filter(self):
             return (ProcedureItem, ModuleItem)
         return ProcedureItem
 
-    def transform_file(self, sourcefile, **kwargs):
-        item = kwargs.get('item', None)
-        if not item and 'items' in kwargs:
-            if kwargs['items']:
-                item = kwargs['items'][0]
-
+    def _get_file_path(self, item, build_args):
         if not item:
             raise ValueError('No Item provided; required to determine file write path')
 
@@ -69,7 +64,26 @@ def transform_file(self, sourcefile, **kwargs):
         path = Path(item.path)
         suffix = self.suffix if self.suffix else path.suffix
         sourcepath = Path(item.path).with_suffix(f'.{_mode}{suffix}')
-        build_args = kwargs.get('build_args', {})
         if build_args and (output_dir := build_args.get('output_dir', None)) is not None:
             sourcepath = Path(output_dir)/sourcepath.name
+        return sourcepath
+
+    def transform_file(self, sourcefile, **kwargs):
+        item = kwargs.get('item')
+        if not item and 'items' in kwargs:
+            if kwargs['items']:
+                item = kwargs['items'][0]
+
+        build_args = kwargs.get('build_args', {})
+        sourcepath = self._get_file_path(item, build_args)
         sourcefile.write(path=sourcepath, cuf=self.cuf)
+
+    def plan_file(self, sourcefile, **kwargs):  # pylint: disable=unused-argument
+        item = kwargs.get('item')
+        if not item and 'items' in kwargs:
+            if kwargs['items']:
+                item = kwargs['items'][0]
+
+        build_args = kwargs.get('build_args', {})
+        sourcepath = self._get_file_path(item, build_args)
+        item.trafo_data['FileWriteTransformation'] = {'path': sourcepath}