#2716 fix tests broken by merge

stfc · Oct 8, 2024 · 0fd9464 · 0fd9464
1 parent b9987fc
commit 0fd9464
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 8 deletions.
diff --git a/src/psyclone/tests/psyir/transformations/kernel_transformation_test.py b/src/psyclone/tests/psyir/transformations/kernel_transformation_test.py
@@ -464,7 +464,8 @@ def test_kernel_gpu_annotation_trans(rtrans, expected_directive,
     rtrans.apply(kern)
 
     # Check that the directive has been added to the kernel code
-    code = fortran_writer(kern.get_kernel_schedule())
+    _, kschedules = kern.get_kernel_schedule()
+    code = fortran_writer(kschedules[0])
     assert expected_directive in code
 
 

diff --git a/src/psyclone/transformations.py b/src/psyclone/transformations.py
@@ -546,15 +546,15 @@ def apply(self, node, options=None):
             node.modified = True
 
             # Get the schedule representing the kernel subroutine
-            routine = node.get_kernel_schedule()
+            _, routines = node.get_kernel_schedule()
         else:
-            routine = node
-
-        for child in routine.children:
-            if isinstance(child, OMPDeclareTargetDirective):
-                return  # The routine is already marked with OMPDeclareTarget
+            routines = [node]
 
-        routine.children.insert(0, OMPDeclareTargetDirective())
+        for routine in routines:
+            for child in routine.children:
+                if isinstance(child, OMPDeclareTargetDirective):
+                    break  # routine is already marked with OMPDeclareTarget
+            routine.children.insert(0, OMPDeclareTargetDirective())
 
     def validate(self, node, options=None):
         ''' Check that an OMPDeclareTargetDirective can be inserted.