(Towards #3449) Support keeping acc routine and omp declare target directives

changelog

@@ -1,6 +1,9 @@
+   30) PR #3451	towards #3449. Adds support for keeping 'ACC routine' and
+   'OMP declare target' directives in existing code.
+
    29) PR #3399 towards #3367. Extends the DefinitionUseChain functionality
    to handle multiple inputs and thus Assignment statements.
- 
+
    28) PR #3417 towards #3398. Move LFRicRedundantComputationTrans into the
    LFRic transformations module.
 

src/psyclone/psyir/frontend/fparser2.py

@@ -60,11 +60,12 @@
 from psyclone.errors import InternalError, GenerationError
 from psyclone.psyir.commentable_mixin import CommentableMixin
 from psyclone.psyir.nodes import (
-    ArrayMember, ArrayOfStructuresReference, ArrayReference, Assignment,
-    BinaryOperation, Call, CodeBlock, Container, DataNode, Directive,
-    FileContainer, IfBlock, IntrinsicCall, Literal, Loop, Member, Node, Range,
-    Reference, Return, Routine, Schedule, StructureReference, UnaryOperation,
-    WhileLoop, Fparser2CodeBlock, ScopingNode, UnknownDirective)
+    ArrayMember, ACCRoutineDirective, ArrayOfStructuresReference,
+    ArrayReference, Assignment, BinaryOperation, Call, CodeBlock, Container,
+    DataNode, Directive, FileContainer, IfBlock, IntrinsicCall, Literal, Loop,
+    Member, Node, OMPDeclareTargetDirective, Range, Reference, Return,
+    Routine, Schedule, StructureReference, UnaryOperation, WhileLoop,
+    Fparser2CodeBlock, ScopingNode, UnknownDirective)
 from psyclone.psyir.nodes.array_mixin import ArrayMixin
 from psyclone.psyir.symbols import (
     ArgumentInterface, ArrayType, AutomaticInterface, ScalarType,
@@ -6145,18 +6146,21 @@ def process_comment(self, comment, preceding_comments):
 
     def _directive_handler(
         self, node: Fortran2003.Directive, parent: Node
-    ) -> Union[CodeBlock, UnknownDirective]:
+    ) -> Union[CodeBlock, UnknownDirective, Directive]:
         '''
         Process a directive and add it to the tree. The current behaviour
         places most directives into a CodeBlock.
 
         Directives starting with !$psy are turned into a UnknownDirective.
 
+        ACC Routine directives and OMP declare target directives are converted
+        to the corresponding PSyIR Directives.
+
         :param node: Directive to process.
         :param parent: The parent to add the PSyIR node to.
 
-        :returns: a CodeBlock containing the input Directive or a
-                  UnknownDirective.
+        :returns: a CodeBlock containing the input Directive, an
+                  UnknownDirective or a specialised PSyIR Directive.
         '''
         # We don't turn OpenMP extensions or directives we can't output
         # correctly into Directive nodes. PSyclone currently always
@@ -6168,6 +6172,30 @@ def _directive_handler(
             not lcase.startswith(prefix) for prefix in dont_match
         ])
         if to_direc:
+            # We first try to specialise some directives.
+            # PSyclone doesn't support clauses on Declare Target so
+            # we can just look for exact string.
+            if lcase == "!$omp declare target":
+                return OMPDeclareTargetDirective(parent=parent)
+
+            if lcase.startswith("!$acc routine"):
+                # If we have an acc routine we need to see if there is
+                # a parallelism clause
+                parallel_clause = lcase[13:].lstrip()
+                if parallel_clause:
+                    try:
+                        directive = ACCRoutineDirective(
+                            parallelism=parallel_clause, parent=parent
+                        )
+                        return directive
+                    except ValueError:
+                        # Fall back to an Unknown Directive if the parallel
+                        # clause isn't understood by PSyclone.
+                        return UnknownDirective(str_rep[2:].lstrip(),
+                                                parent=parent)
+                # If we have no parallel clause then return the default.
+                return ACCRoutineDirective(parent=parent)
+            # Otherwise return an UnknownDirective for the node.
             content = str_rep[2:].lstrip()
             return UnknownDirective(content, parent=parent)
         code_block = Fparser2CodeBlock(

src/psyclone/tests/psyir/frontend/fparser2_directive_test.py

@@ -40,7 +40,8 @@
 
 from psyclone.psyir.frontend.fortran import FortranReader
 from psyclone.psyir.nodes import (
-    CodeBlock, IfBlock, UnknownDirective
+    ACCRoutineDirective, CodeBlock, IfBlock, OMPDeclareTargetDirective,
+    UnknownDirective
 )
 
 
@@ -457,3 +458,72 @@ def test_comments_on_directive_before_where(fortran_writer):
     end if
   enddo"""
     assert correct in fortran_writer(psyir)
+
+
+def test_omp_declare_target_directive(fortran_writer):
+    '''Tests that an omp declare target directive is converted to the
+    corresponding directive node.'''
+    code = """subroutine x
+    real, dimension(100) :: a, b
+    !$OMP declare target
+    a = b + 1
+    end subroutine x"""
+
+    reader = FortranReader(ignore_comments=False, ignore_directives=False)
+    psyir = reader.psyir_from_source(code)
+    routine = psyir.children[0]
+    assert isinstance(routine.children[0], OMPDeclareTargetDirective)
+    correct = """  !$omp declare target
+  a = b + 1"""
+    assert correct in fortran_writer(psyir)
+
+
+def test_acc_routine_directive(fortran_writer):
+    '''Tests that acc routine directives give the corresponding directive
+    node unless the parallel declaration isn't understood.'''
+    reader = FortranReader(ignore_comments=False, ignore_directives=False)
+    # Test with undeclared parallel type.
+    code = """subroutine x
+    real, dimension(100) :: a, b
+    !$acc routine
+    a = b + 1
+    end subroutine x"""
+
+    psyir = reader.psyir_from_source(code)
+    routine = psyir.children[0]
+    assert isinstance(routine.children[0], ACCRoutineDirective)
+    assert routine.children[0].parallelism == "seq"
+
+    correct = """  !$acc routine seq
+  a = b + 1"""
+    assert correct in fortran_writer(psyir)
+
+    # Test with declared parallel type.
+    code = """subroutine x
+    real, dimension(100) :: a, b
+    !$acc routine vector
+    a = b + 1
+    end subroutine x"""
+
+    psyir = reader.psyir_from_source(code)
+    routine = psyir.children[0]
+    assert isinstance(routine.children[0], ACCRoutineDirective)
+    assert routine.children[0].parallelism == "vector"
+
+    correct = """  !$acc routine vector
+  a = b + 1"""
+    assert correct in fortran_writer(psyir)
+
+    # Test with declared parallel type.
+    code = """subroutine x
+    real, dimension(100) :: a, b
+    !$acc routine unknown
+    a = b + 1
+    end subroutine x"""
+
+    psyir = reader.psyir_from_source(code)
+    routine = psyir.children[0]
+    assert isinstance(routine.children[0], UnknownDirective)
+    correct = """  !$acc routine unknown
+  a = b + 1"""
+    assert correct in fortran_writer(psyir)