From 0fd946484f9ecbdfecc415e54baced1e393a2b1b Mon Sep 17 00:00:00 2001 From: Andrew Porter Date: Tue, 8 Oct 2024 15:08:36 +0100 Subject: [PATCH] #2716 fix tests broken by merge --- .../transformations/kernel_transformation_test.py | 3 ++- src/psyclone/transformations.py | 14 +++++++------- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/src/psyclone/tests/psyir/transformations/kernel_transformation_test.py b/src/psyclone/tests/psyir/transformations/kernel_transformation_test.py index 14a7b19a01..a6b0bebe44 100644 --- 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 index 9904a36fad..b3913d11ed 100644 --- 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.