From 8992b755fe53848f2b9078190a782ce9c0c427e2 Mon Sep 17 00:00:00 2001
From: Sergi Siso <sergiesg@gmail.com>
Date: Wed, 2 Oct 2024 14:58:57 +0100
Subject: [PATCH] #2730 Add gen_code support for OpenMP target directives

---
 examples/lfric/scripts/gpu_offloading.py   |  4 ++--
 src/psyclone/f2pygen.py                    |  2 +-
 src/psyclone/psyir/nodes/omp_directives.py | 13 +++++++++++++
 src/psyclone/transformations.py            | 15 +++++++++++++--
 4 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/examples/lfric/scripts/gpu_offloading.py b/examples/lfric/scripts/gpu_offloading.py
index 67706e89e6..b665dfcf94 100644
--- a/examples/lfric/scripts/gpu_offloading.py
+++ b/examples/lfric/scripts/gpu_offloading.py
@@ -137,8 +137,8 @@ def trans(psy):
         # Add GPU offloading to loops unless they are over colours or are null.
         schedule = invoke.schedule
         for loop in schedule.walk(Loop):
-            if offload or all(kern.name.lower() in failed_to_offload for
-                              kern in loop.kernels()):
+            if offload and all(kern.name.lower() not in failed_to_offload for
+                               kern in loop.kernels()):
                 try:
                     if loop.loop_type == "colours":
                         pass
diff --git a/src/psyclone/f2pygen.py b/src/psyclone/f2pygen.py
index 4db3407e44..bfe41981ec 100644
--- a/src/psyclone/f2pygen.py
+++ b/src/psyclone/f2pygen.py
@@ -138,7 +138,7 @@ class OMPDirective(Directive):
     '''
     def __init__(self, root, line, position, dir_type):
         self._types = ["parallel do", "parallel", "do", "master", "single",
-                       "taskloop", "taskwait", "declare"]
+                       "taskloop", "taskwait", "declare", "target"]
         self._positions = ["begin", "end"]
 
         super(OMPDirective, self).__init__(root, line, position, dir_type)
diff --git a/src/psyclone/psyir/nodes/omp_directives.py b/src/psyclone/psyir/nodes/omp_directives.py
index 00318273ca..722c67b032 100644
--- a/src/psyclone/psyir/nodes/omp_directives.py
+++ b/src/psyclone/psyir/nodes/omp_directives.py
@@ -2402,6 +2402,19 @@ def end_string(self):
         '''
         return "omp end target"
 
+    def gen_code(self, parent):
+        '''Generate the OpenMP Target Directive and any associated code.
+
+        :param parent: the parent Node in the Schedule to which to add our
+                       content.
+        :type parent: sub-class of :py:class:`psyclone.f2pygen.BaseGen`
+        '''
+        # Check the constraints are correct
+        self.validate_global_constraints()
+
+        # Generate the code for this Directive
+        parent.add(DirectiveGen(parent, "omp", "begin", "target"))
+
 
 class OMPLoopDirective(OMPRegionDirective):
     ''' Class for the !$OMP LOOP directive that specifies that the iterations
diff --git a/src/psyclone/transformations.py b/src/psyclone/transformations.py
index ca34b76607..2d20acf43c 100644
--- a/src/psyclone/transformations.py
+++ b/src/psyclone/transformations.py
@@ -546,10 +546,21 @@ def apply(self, node, options=None):
 
         '''
         self.validate(node, options)
-        for child in node.children:
+
+        if isinstance(node, Kern):
+            # Flag that the kernel has been modified
+            node.modified = True
+
+            # Get the schedule representing the kernel subroutine
+            routine = node.get_kernel_schedule()
+        else:
+            routine = node
+
+        for child in routine.children:
             if isinstance(child, OMPDeclareTargetDirective):
                 return  # The routine is already marked with OMPDeclareTarget
-        node.children.insert(0, OMPDeclareTargetDirective())
+
+        routine.children.insert(0, OMPDeclareTargetDirective())
 
     def validate(self, node, options=None):
         ''' Check that an OMPDeclareTargetDirective can be inserted.