Source code for psyclone.psyir.transformations.omp_declare_target_trans

# -----------------------------------------------------------------------------
# SPDX-FileCopyrightText: Copyright (c) 2017-2026 Science and Technology
#                         Facilities Council
# SPDX-License-Identifier: BSD-3-Clause
# See the full LICENSE file in the project root for details.
# -----------------------------------------------------------------------------

'''
This module provides the implementation of OMPDeclareTargetTrans

'''
from typing import Any, Optional, Union

from psyclone.psyir.nodes import OMPDeclareTargetDirective, Routine
from psyclone.psyir.transformations.callee_transformation_mixin import (
    CalleeTransformationMixin)
from psyclone.psyGen import Transformation, Kern
from psyclone.psyir.transformations.mark_routine_for_gpu_mixin import (
    MarkRoutineForGPUMixin)
from psyclone.utils import transformation_documentation_wrapper




@transformation_documentation_wrapper
class OMPDeclareTargetTrans(Transformation,
                            MarkRoutineForGPUMixin,
                            CalleeTransformationMixin):
    '''
    Adds an OpenMP declare target directive to the specified routine.

    For example:

    >>> from psyclone.psyir.frontend.fortran import FortranReader
    >>> from psyclone.psyir.backend.fortran import FortranWriter
    >>> from psyclone.psyir.nodes import Loop
    >>> from psyclone.psyir.transformations import OMPDeclareTargetTrans
    >>>
    >>> tree = FortranReader().psyir_from_source("""
    ...     subroutine my_subroutine(A)
    ...         integer, dimension(10, 10), intent(inout) :: A
    ...         integer :: i
    ...         integer :: j
    ...         do i = 1, 10
    ...             do j = 1, 10
    ...                 A(i, j) = 0
    ...             end do
    ...         end do
    ...     end subroutine
    ...     """)
    >>> omptargettrans = OMPDeclareTargetTrans()
    >>> omptargettrans.apply(tree.walk(Routine)[0])
    >>> print(FortranWriter()(tree))
    subroutine my_subroutine(a)
      integer, dimension(10,10), intent(inout) :: a
      integer :: i
      integer :: j
    <BLANKLINE>
      !$omp declare target
      do i = 1, 10, 1
        do j = 1, 10, 1
          a(i,j) = 0
        enddo
      enddo
    <BLANKLINE>
    end subroutine my_subroutine
    <BLANKLINE>

    '''


    def apply(self,
              node: Union[Routine, Kern],
              options: Optional[dict[str, Any]] = None,
              force: bool = False,
              device_string: str = "",
              **kwargs) -> None:
        ''' Insert an OMPDeclareTargetDirective inside the provided routine or
        associated PSyKAl kernel.

        :param node: the kernel or routine which is the target of this
                     transformation.
        :param options: a dictionary with options for transformations.
        :param force: whether to allow routines with CodeBlocks to run on
                      the GPU.
        :param device_string: provide a compiler-platform identifier.

        '''
        self.validate(node, options, force=force, device_string=device_string,
                      **kwargs)

        if isinstance(node, Kern):
            # Get the schedule representing the kernel subroutine
            routines = node.get_callees()
        else:
            routines = [node]

        for routine in routines:
            if not any(isinstance(child, OMPDeclareTargetDirective) for
                       child in routine.children):
                routine.children.insert(0, OMPDeclareTargetDirective())




    def validate(self,
                 node: Union[Kern, Routine],
                 options: Optional[dict[str, Any]] = None,
                 **kwargs) -> None:
        ''' Check that an OMPDeclareTargetDirective can be inserted.

        :param node: the kernel or routine which is the target of this
            transformation.
        :param options: a dictionary with options for transformations.

        '''
        # TODO #2668 Depracate options dict.
        if not options:
            self.validate_options(**kwargs)
        super().validate(node, options=options, **kwargs)

        self.validate_it_can_run_on_gpu(node, options, **kwargs)

        if isinstance(node, Kern):
            self._check_callee_implementation_is_local(node)