Scheduler.process: apply directly to program unit corresponding to Item

loki/bulk/scheduler.py

@@ -599,7 +599,8 @@ def item_successors(self, item):
                 successors += [self.item_map[child.name]] + self.item_successors(child)
         return successors
 
-    def process(self, transformation, reverse=False, item_filter=SubroutineItem, use_file_graph=False):
+    def process(self, transformation, reverse=False, item_filter=SubroutineItem, use_file_graph=False,
+                recurse_to_contained_nodes=False):
         """
         Process all :attr:`items` in the scheduler's graph
 
@@ -612,6 +613,12 @@ def process(self, transformation, reverse=False, item_filter=SubroutineItem, use
         by setting :data:`use_file_graph` to ``True``. Currently, this calls
         the transformation on the first `item` associated with a file only.
         In this mode, :data:`item_filter` does not have any effect.
+
+        The scheduler applies the transformation to the IR node corresponding to
+        each item in the scheduler's graph. Optionally, the transformation can also
+        be applied recursively, e.g., to all member subroutines contained in a
+        subroutine object, or all modules and subroutines contained in a source file,
+        by setting :data:`recurse_to_contained_nodes` to ``True``.
         """
         trafo_name = transformation.__class__.__name__
         log = f'[Loki::Scheduler] Applied transformation <{trafo_name}>' + ' in {:.2f}s'
@@ -629,7 +636,14 @@ def process(self, transformation, reverse=False, item_filter=SubroutineItem, use
             if use_file_graph:
                 for node in traversal:
                     items = graph.nodes[node]['items']
-                    transformation.apply(items[0].source, item=items[0], items=items, recurse_to_contained_nodes=True)
+
+                    if item_filter and any(not isinstance(item, item_filter) for item in items):
+                        continue
+
+                    transformation.apply(
+                        items[0].source, item=items[0], items=items,
+                        recurse_to_contained_nodes=recurse_to_contained_nodes
+                    )
             else:
                 for item in traversal:
                     if item_filter and not isinstance(item, item_filter):
@@ -638,12 +652,19 @@ def process(self, transformation, reverse=False, item_filter=SubroutineItem, use
                     # Use entry from item_map to ensure the original item is used in transformation
                     _item = self.item_map[item.name]
 
+                    # Pick out the IR node to which to apply the transformation
+                    # TODO: should this become an Item property?
+                    if isinstance(item, SubroutineItem):
+                        source = _item.routine
+                    else:
+                        source = _item.scope
+
                     # Process work item with appropriate kernel
                     transformation.apply(
-                        _item.source, role=_item.role, mode=_item.mode,
+                        source, role=_item.role, mode=_item.mode,
                         item=_item, targets=_item.targets,
                         successors=self.item_successors(_item), depths=self.depths,
-                        recurse_to_contained_nodes=True
+                        recurse_to_contained_nodes=recurse_to_contained_nodes
                     )
 
     def callgraph(self, path, with_file_graph=False):

loki/transform/dependency_transform.py

@@ -79,11 +79,35 @@ def transform_subroutine(self, routine, **kwargs):
         """
         role = kwargs.get('role')
 
-        if role == 'kernel':
+        # If applied recursively over all routines in a sourcefile, we may
+        # visit contained member subroutines. Since these are not outward facing
+        # it is not necessary to update their name.
+        is_member = isinstance(routine.parent, Subroutine)
+
+        # Pick out correct item for current subroutine if processed via recursion
+        if 'items' in kwargs:
+            item_names = [item_.local_name for item_ in kwargs['items']]
+            if routine.name.lower() in item_names:
+                kwargs['item'] = kwargs['items'][item_names.index(routine.name.lower())]
+
+        # If called without explicit role or target, extract from Item
+        # We need to do this here to cache the value for targets, because
+        # rename_calls will change this property
+        if (item := kwargs.get('item')):
+            if not role:
+                kwargs['role'] = item.role
+                role = item.role
+            if not kwargs.get('targets'):
+                kwargs['targets'] = item.targets
+
+        if role == 'kernel' and not is_member:
+            if routine.name.endswith(self.suffix):
+                # This is to ensure that the transformation is idempotent if
+                # applied more than once to a routine
+                return
             # Change the name of kernel routines
-            if routine.is_function:
-                if not routine.result_name:
-                    self.update_result_var(routine)
+            if routine.is_function and not routine.result_name:
+                self.update_result_var(routine)
             routine.name += self.suffix
 
         self.rename_calls(routine, **kwargs)
@@ -96,7 +120,7 @@ def transform_subroutine(self, routine, **kwargs):
         intfs = FindNodes(Interface).visit(routine.spec)
         self.rename_interfaces(routine, intfs=intfs, **kwargs)
 
-        if role == 'kernel' and self.mode == 'strict':
+        if role == 'kernel' and self.mode == 'strict' and not is_member:
             # Re-generate C-style interface header
             self.generate_interfaces(routine)
 
@@ -120,6 +144,8 @@ def transform_module(self, module, **kwargs):
         Rename kernel modules and re-point module-level imports.
         """
         role = kwargs.get('role')
+        if not role and 'item' in kwargs:
+            role = kwargs['item'].role
 
         if role == 'kernel':
             # Change the name of kernel modules
@@ -134,6 +160,8 @@ def transform_file(self, sourcefile, **kwargs):
         In 'module' mode perform module-wrapping for dependnecy injection.
         """
         role = kwargs.get('role')
+        if not role and 'item' in kwargs:
+            role = kwargs['item'].role
         if role == 'kernel' and self.mode == 'module':
             self.module_wrap(sourcefile, **kwargs)
 
@@ -144,9 +172,9 @@ def rename_calls(self, routine, **kwargs):
         :param targets: Optional list of subroutine names for which to
                         modify the corresponding calls.
         """
-        targets = as_tuple(kwargs.get('targets', None))
+        targets = as_tuple(kwargs.get('targets'))
         targets = as_tuple(str(t).upper() for t in targets)
-        members = [r.name for r in routine.subroutines]
+        members = [r.name.upper() for r in routine.subroutines]
 
         if self.replace_ignore_items:
             item = kwargs.get('item', None)
@@ -159,6 +187,8 @@ def rename_calls(self, routine, **kwargs):
                 call._update(name=call.name.clone(name=f'{call.name}{self.suffix}'))
 
         for call in FindInlineCalls(unique=False).visit(routine.body):
+            if call.name.upper() in members:
+                continue
             if targets is None or call.name.upper() in targets:
                 call.function = call.function.clone(name=f'{call.name}{self.suffix}')
 
@@ -169,7 +199,9 @@ def rename_imports(self, source, imports, **kwargs):
         :param targets: Optional list of subroutine names for which to
                         modify the corresponding calls.
         """
-        targets = as_tuple(kwargs.get('targets', None))
+        targets = as_tuple(kwargs.get('targets'))
+        if not targets and 'item' in kwargs:
+            targets = as_tuple(kwargs['item'].targets)
         targets = as_tuple(str(t).upper() for t in targets)
 
         # We don't want to rename module variable imports, so we build
@@ -234,10 +266,12 @@ def rename_interfaces(self, source, intfs, **kwargs):
         Update explicit interfaces to actively transformed subroutines.
         """
 
-        targets = as_tuple(kwargs.get('targets', None))
+        targets = as_tuple(kwargs.get('targets'))
+        if not targets and 'item' in kwargs:
+            targets = as_tuple(kwargs['item'].targets)
         targets = as_tuple(str(t).lower() for t in targets)
 
-        if self.replace_ignore_items and (item := kwargs.get('item', None)):
+        if self.replace_ignore_items and (item := kwargs.get('item')):
             targets += as_tuple(str(i).lower() for i in item.ignore)
 
         # Transformer map to remove any outdated interfaces
@@ -296,6 +330,8 @@ def module_wrap(self, sourcefile, **kwargs):
         Wrap target subroutines in modules and replace in source file.
         """
         targets = as_tuple(kwargs.get('targets', None))
+        if not targets and 'item' in kwargs:
+            targets = as_tuple(kwargs['item'].targets)
         targets = as_tuple(str(t).upper() for t in targets)
         item = kwargs.get('item', None)
 
@@ -305,7 +341,7 @@ def module_wrap(self, sourcefile, **kwargs):
         for routine in sourcefile.subroutines:
             if routine not in module_routines:
                 # Skip member functions
-                if item and f'#{routine.name.lower()}' != item.name.lower():
+                if item and routine.name.lower() != item.local_name.lower():
                     continue
 
                 # Skip internal utility routines too

loki/transform/transformation.py

@@ -173,10 +173,6 @@ def apply_subroutine(self, subroutine, recurse_to_contained_nodes=False, **kwarg
         if subroutine._incomplete:
             raise RuntimeError('Transformation.apply_subroutine requires Subroutine to be complete')
 
-        # Bail if the subroutine has not actually been scheduled for processing
-        if (item := kwargs.get('item', None)) and item.local_name != subroutine.name.lower():
-            return
-
         # Apply the actual transformation for subroutines
         self.transform_subroutine(subroutine, **kwargs)
 

scripts/loki_transform.py

@@ -300,11 +300,17 @@ def transform_subroutine(self, routine, **kwargs):
     mode = mode.replace('-', '_')  # Sanitize mode string
     dependency = DependencyTransformation(suffix=f'_{mode.upper()}',
                                           mode='module', module_suffix='_MOD')
-    scheduler.process(transformation=dependency)
+    scheduler.process(transformation=dependency, use_file_graph=True, recurse_to_contained_nodes=True)
 
     # Write out all modified source files into the build directory
-    scheduler.process(transformation=FileWriteTransformation(builddir=out_path, mode=mode, cuf='cuf' in mode),
-                      use_file_graph=True)
+    if global_var_offload:
+        item_filter = (SubroutineItem, GlobalVarImportItem)
+    else:
+        item_filter = SubroutineItem
+    scheduler.process(
+        transformation=FileWriteTransformation(builddir=out_path, mode=mode, cuf='cuf' in mode),
+        use_file_graph=True, item_filter=item_filter
+    )
 
 
 @cli.command()
@@ -538,10 +544,10 @@ def ecphys(mode, config, header, source, build, cpp, directive, frontend):
     dependency = DependencyTransformation(
         mode='module', module_suffix='_MOD', suffix=f'_{mode.upper()}'
     )
-    scheduler.process(transformation=dependency)
+    scheduler.process(transformation=dependency, use_file_graph=True, recurse_to_contained_nodes=True)
 
     # Write out all modified source files into the build directory
-    scheduler.process(transformation=FileWriteTransformation(builddir=build, mode=mode))
+    scheduler.process(transformation=FileWriteTransformation(builddir=build, mode=mode), use_file_graph=True)
 
 
 if __name__ == "__main__":