Merge pull request #111 from ecmwf-ifs/naml-remove-call-kernel-only

ECPhys: Add a "kernel-only" option to remove calls

scripts/loki_transform.py

@@ -468,7 +468,7 @@ def ecphys(mode, config, header, source, build, cpp, directive, frontend):
     # Remove DR_HOOK and other utility calls first, so they don't interfere with SCC loop hoisting
     if 'scc' in mode:
         scheduler.process(transformation=RemoveCallsTransformation(
-            routines=['DR_HOOK', 'ABOR1', 'WRITE(NULOUT'], include_intrinsics=True
+            routines=['DR_HOOK', 'ABOR1', 'WRITE(NULOUT'], include_intrinsics=True, kernel_only=True
         ))
     else:
         scheduler.process(transformation=DrHookTransformation(mode=mode, remove=False))

transformations/tests/test_utility_routines.py

@@ -171,18 +171,20 @@ def test_dr_hook_transformation_remove(frontend, config, source):
 
 
 @pytest.mark.parametrize('include_intrinsics', (True, False))
+@pytest.mark.parametrize('kernel_only', (True, False))
 @pytest.mark.parametrize('frontend', available_frontends(
     xfail=[(OMNI, 'Incomplete source tree impossible with OMNI')]
 ))
-def test_utility_routine_removal(frontend, config, source, include_intrinsics):
+def test_utility_routine_removal(frontend, config, source, include_intrinsics, kernel_only):
     """
     Test removal of utility calls and intrinsics with custom patterns.
     """
     scheduler_config = SchedulerConfig.from_dict(config)
     scheduler = Scheduler(paths=source, config=scheduler_config, frontend=frontend)
     scheduler.process(
         transformation=RemoveCallsTransformation(
-            routines=['ABOR1', 'WRITE(NULOUT', 'DR_HOOK'], include_intrinsics=include_intrinsics
+            routines=['ABOR1', 'WRITE(NULOUT', 'DR_HOOK'],
+            include_intrinsics=include_intrinsics, kernel_only=kernel_only
         )
     )
 
@@ -197,3 +199,7 @@ def test_utility_routine_removal(frontend, config, source, include_intrinsics):
     routine = Sourcefile.from_file(source/'never_gonna_give.F90', frontend=frontend)['i_hope_you_havent_let_me_down']
     assert 'zhook_handle' in routine.variables
     assert len([call for call in FindNodes(CallStatement).visit(routine.body) if call.name == 'dr_hook']) == 2
+
+    driver = scheduler.item_map['#rick_astley'].routine
+    drhook_calls = [call for call in FindNodes(CallStatement).visit(driver.body) if call.name == 'dr_hook']
+    assert len(drhook_calls) == (2 if kernel_only else 0)

transformations/transformations/utility_routines.py

@@ -103,16 +103,25 @@ class RemoveCallsTransformation(Transformation):
     include_intrinsics : bool
         Option to extend searches to :any:`Intrinsic` nodes to
         capture print/write statements
+    kernel_only : bool
+        Option to only remove calls in routines marked as "kernel"; default: ``False``
     """
-    def __init__(self, routines, include_intrinsics=False, **kwargs):
+    def __init__(self, routines, include_intrinsics=False, kernel_only=False, **kwargs):
         self.routines = as_tuple(routines)
         self.include_intrinsics = include_intrinsics
+        self.kernel_only = kernel_only
         super().__init__(**kwargs)
 
     def transform_subroutine(self, routine, **kwargs):
         """
         Apply transformation to subroutine object
         """
+
+        # Skip driver layer if requested
+        role = kwargs.get('role', None)
+        if role and role == 'driver' and self.kernel_only:
+            return
+
         mapper = {}
 
         # First remove inline conditionals with calls to specified routines or intrinsics