diff --git a/src/psyclone/domain/lfric/__init__.py b/src/psyclone/domain/lfric/__init__.py index 4130e201fa..0076185d0a 100644 --- a/src/psyclone/domain/lfric/__init__.py +++ b/src/psyclone/domain/lfric/__init__.py @@ -77,7 +77,8 @@ from psyclone.domain.lfric.lfric_invoke_schedule import LFRicInvokeSchedule from psyclone.domain.lfric.lfric_dofmaps import LFRicDofmaps from psyclone.domain.lfric.lfric_stencils import LFRicStencils - +from psyclone.domain.lfric.lfric_omp_parallel_loop_trans import \ + LFRicOMPParallelLoopTrans __all__ = [ 'ArgOrdering', diff --git a/src/psyclone/domain/lfric/lfric_async_halo_exchange_trans.py b/src/psyclone/domain/lfric/lfric_async_halo_exchange_trans.py new file mode 100644 index 0000000000..d53e51d07d --- /dev/null +++ b/src/psyclone/domain/lfric/lfric_async_halo_exchange_trans.py @@ -0,0 +1,131 @@ +# ----------------------------------------------------------------------------- +# BSD 3-Clause License +# +# Copyright (c) 2017-2025, Science and Technology Facilities Council. +# All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# * Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# * Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# * Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS +# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE +# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT +# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN +# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +# POSSIBILITY OF SUCH DAMAGE. +# ----------------------------------------------------------------------------- +# Authors R. W. Ford, A. R. Porter, S. Siso and N. Nobre, STFC Daresbury Lab +# A. B. G. Chalk and V. K. Atkinson, STFC Daresbury Lab +# J. Henrichs, Bureau of Meteorology +# Modified I. Kavcic, J. G. Wallwork, O. Brunt and L. Turner, Met Office +# S. Valat, Inria / Laboratoire Jean Kuntzmann +# M. Schreiber, Univ. Grenoble Alpes / Inria / Lab. Jean Kuntzmann +# J. Dendy, Met Office + +from psyclone.psyGen import Transformation +from psyclone.lfric import LFRicHaloExchangeEnd, LFRicHaloExchangeStart +from psyclone import psyGen +from psyclone.psyir.transformations.transformation_error import ( + TransformationError) + + +class LFRicAsyncHaloExchangeTrans(Transformation): + '''Splits a synchronous halo exchange into a halo exchange start and + halo exchange end. For example: + + >>> from psyclone.parse.algorithm import parse + >>> from psyclone.psyGen import PSyFactory + >>> api = "lfric" + >>> ast, invokeInfo = parse("file.f90", api=api) + >>> psy=PSyFactory(api).create(invokeInfo) + >>> schedule = psy.invokes.get('invoke_0').schedule + >>> # Uncomment the following line to see a text view of the schedule + >>> # print(schedule.view()) + >>> + >>> from psyclone.transformations import LFRicAsyncHaloExchangeTrans + >>> trans = LFRicAsyncHaloExchangeTrans() + >>> trans.apply(schedule.children[0]) + >>> # Uncomment the following line to see a text view of the schedule + >>> # print(schedule.view()) + + ''' + + def __str__(self): + return "Changes a synchronous halo exchange into an asynchronous one." + + @property + def name(self): + ''' + :returns: the name of this transformation as a string. + :rtype: str + ''' + return "LFRicAsyncHaloExchangeTrans" + + def apply(self, node, options=None): + '''Transforms a synchronous halo exchange, represented by a + HaloExchange node, into an asynchronous halo exchange, + represented by HaloExchangeStart and HaloExchangeEnd nodes. + + :param node: a synchronous haloexchange node. + :type node: :py:obj:`psyclone.psygen.HaloExchange` + :param options: a dictionary with options for transformations. + :type options: Optional[Dict[str, Any]] + + ''' + self.validate(node, options) + + # add asynchronous start and end halo exchanges and initialise + # them using information from the existing synchronous halo + # exchange + # pylint: disable=protected-access + node.parent.addchild( + LFRicHaloExchangeStart( + node.field, check_dirty=node._check_dirty, + vector_index=node.vector_index, parent=node.parent), + index=node.position) + node.parent.addchild( + LFRicHaloExchangeEnd( + node.field, check_dirty=node._check_dirty, + vector_index=node.vector_index, parent=node.parent), + index=node.position) + + # remove the existing synchronous halo exchange + node.detach() + + def validate(self, node, options): + # pylint: disable=signature-differs + '''Internal method to check whether the node is valid for this + transformation. + + :param node: a synchronous Halo Exchange node + :type node: :py:obj:`psyclone.psygen.HaloExchange` + :param options: a dictionary with options for transformations. + :type options: Optional[Dict[str, Any]] + + :raises TransformationError: if the node argument is not a + HaloExchange (or subclass thereof) + + ''' + if not isinstance(node, psyGen.HaloExchange) or \ + isinstance(node, (LFRicHaloExchangeStart, LFRicHaloExchangeEnd)): + raise TransformationError( + f"Error in LFRicAsyncHaloExchange transformation. Supplied" + f" node must be a synchronous halo exchange but found " + f"'{type(node)}'.") diff --git a/src/psyclone/domain/lfric/lfric_omp_parallel_loop_trans.py b/src/psyclone/domain/lfric/lfric_omp_parallel_loop_trans.py new file mode 100644 index 0000000000..2344814a24 --- /dev/null +++ b/src/psyclone/domain/lfric/lfric_omp_parallel_loop_trans.py @@ -0,0 +1,107 @@ +# ----------------------------------------------------------------------------- +# BSD 3-Clause License +# +# Copyright (c) 2017-2025, Science and Technology Facilities Council. +# All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# * Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# * Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# * Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS +# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE +# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT +# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN +# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +# POSSIBILITY OF SUCH DAMAGE. +# ----------------------------------------------------------------------------- +# Authors R. W. Ford, A. R. Porter, S. Siso and N. Nobre, STFC Daresbury Lab +# A. B. G. Chalk and V. K. Atkinson STFC Daresbury Lab +# J. Henrichs, Bureau of Meteorology +# Modified I. Kavcic, J. G. Wallwork, O. Brunt and L. Turner, Met Office +# S. Valat, Inria / Laboratoire Jean Kuntzmann +# M. Schreiber, Univ. Grenoble Alpes / Inria / Lab. Jean Kuntzmann +# J. Dendy, Met Office + +from psyclone.psyir.transformations.omp_parallel_loop_trans import ( + OMPParallelLoopTrans) +from psyclone.domain.lfric import (LFRicConstants, LFRicLoop) +from psyclone.psyir.transformations.transformation_error import ( + TransformationError) + + +class LFRicOMPParallelLoopTrans(OMPParallelLoopTrans): + + ''' LFRic-specific OpenMP loop transformation. Adds LFRic specific + validity checks. Actual transformation is done by the + :py:class:`base class `. + + :param str omp_directive: choose which OpenMP loop directive to use. + Defaults to "do". + :param str omp_schedule: the OpenMP schedule to use. Must be one of + 'runtime', 'static', 'dynamic', 'guided' or 'auto'. Defaults to + 'static'. + + ''' + def __init__(self, omp_directive="do", omp_schedule="static"): + super().__init__(omp_directive=omp_directive, + omp_schedule=omp_schedule) + + def __str__(self): + return "Add an OpenMP Parallel Do directive to an LFRic loop" + + def validate(self, node, options=None): + ''' + Perform LFRic-specific loop validity checks then call the `validate` + method of the base class. + + :param node: the Node in the Schedule to check + :type node: :py:class:`psyclone.psyir.nodes.Node` + :param options: a dictionary with options for transformations. + :type options: Optional[Dict[str, Any]] + + :raises TransformationError: if the supplied Node is not an LFRicLoop. + :raises TransformationError: if the associated loop requires + colouring. + ''' + if not isinstance(node, LFRicLoop): + raise TransformationError( + f"Error in {self.name} transformation. The supplied node " + f"must be an LFRicLoop but got '{type(node).__name__}'") + + # If the loop is not already coloured then check whether or not + # it should be. If the field space is discontinuous (including + # any_discontinuous_space) then we don't need to worry about + # colouring. + const = LFRicConstants() + if node.field_space.orig_name not in const.VALID_DISCONTINUOUS_NAMES: + if (node.loop_type not in ('cells_in_colour', 'tiles_in_colour') + and node.has_inc_arg()): + raise TransformationError( + f"Error in {self.name} transformation. The kernel has an " + f"argument with INC access but the loop is of type " + f"'{node.loop_type}'. Colouring is required.") + # As this is a domain-specific loop, we don't perform general + # dependence analysis because it is too conservative and doesn't + # account for the special steps taken for such a loop at code- + # generation time (e.g. the way we ensure variables are given the + # correct sharing attributes). + local_options = options.copy() if options else {} + local_options["force"] = True + super().validate(node, options=local_options) diff --git a/src/psyclone/transformations.py b/src/psyclone/transformations.py index 45359f56a6..373509890b 100644 --- a/src/psyclone/transformations.py +++ b/src/psyclone/transformations.py @@ -53,7 +53,6 @@ from psyclone.core import Signature, VariablesAccessMap from psyclone.domain.lfric import (KernCallArgList, LFRicConstants, LFRicInvokeSchedule, LFRicKern, LFRicLoop) -from psyclone.lfric import LFRicHaloExchangeEnd, LFRicHaloExchangeStart from psyclone.errors import InternalError from psyclone.gocean1p0 import GOInvokeSchedule from psyclone.psyGen import (Transformation, CodedKern, Kern, InvokeSchedule) @@ -112,67 +111,6 @@ def check_intergrid(node): f" is such a kernel.") -class LFRicOMPParallelLoopTrans(OMPParallelLoopTrans): - - ''' LFRic-specific OpenMP loop transformation. Adds LFRic specific - validity checks. Actual transformation is done by the - :py:class:`base class `. - - :param str omp_directive: choose which OpenMP loop directive to use. - Defaults to "do". - :param str omp_schedule: the OpenMP schedule to use. Must be one of - 'runtime', 'static', 'dynamic', 'guided' or 'auto'. Defaults to - 'static'. - - ''' - def __init__(self, omp_directive="do", omp_schedule="static"): - super().__init__(omp_directive=omp_directive, - omp_schedule=omp_schedule) - - def __str__(self): - return "Add an OpenMP Parallel Do directive to an LFRic loop" - - def validate(self, node, options=None): - ''' - Perform LFRic-specific loop validity checks then call the `validate` - method of the base class. - - :param node: the Node in the Schedule to check - :type node: :py:class:`psyclone.psyir.nodes.Node` - :param options: a dictionary with options for transformations. - :type options: Optional[Dict[str, Any]] - - :raises TransformationError: if the supplied Node is not an LFRicLoop. - :raises TransformationError: if the associated loop requires - colouring. - ''' - if not isinstance(node, LFRicLoop): - raise TransformationError( - f"Error in {self.name} transformation. The supplied node " - f"must be an LFRicLoop but got '{type(node).__name__}'") - - # If the loop is not already coloured then check whether or not - # it should be. If the field space is discontinuous (including - # any_discontinuous_space) then we don't need to worry about - # colouring. - const = LFRicConstants() - if node.field_space.orig_name not in const.VALID_DISCONTINUOUS_NAMES: - if (node.loop_type not in ('cells_in_colour', 'tiles_in_colour') - and node.has_inc_arg()): - raise TransformationError( - f"Error in {self.name} transformation. The kernel has an " - f"argument with INC access but the loop is of type " - f"'{node.loop_type}'. Colouring is required.") - # As this is a domain-specific loop, we don't perform general - # dependence analysis because it is too conservative and doesn't - # account for the special steps taken for such a loop at code- - # generation time (e.g. the way we ensure variables are given the - # correct sharing attributes). - local_options = options.copy() if options else {} - local_options["force"] = True - super().validate(node, options=local_options) - - class LFRicOMPLoopTrans(OMPLoopTrans): ''' LFRic specific orphan OpenMP loop transformation. Adds @@ -1323,91 +1261,6 @@ def apply(self, loop, options=None): loop.update_halo_exchanges() -class LFRicAsyncHaloExchangeTrans(Transformation): - '''Splits a synchronous halo exchange into a halo exchange start and - halo exchange end. For example: - - >>> from psyclone.parse.algorithm import parse - >>> from psyclone.psyGen import PSyFactory - >>> api = "lfric" - >>> ast, invokeInfo = parse("file.f90", api=api) - >>> psy=PSyFactory(api).create(invokeInfo) - >>> schedule = psy.invokes.get('invoke_0').schedule - >>> # Uncomment the following line to see a text view of the schedule - >>> # print(schedule.view()) - >>> - >>> from psyclone.transformations import LFRicAsyncHaloExchangeTrans - >>> trans = LFRicAsyncHaloExchangeTrans() - >>> trans.apply(schedule.children[0]) - >>> # Uncomment the following line to see a text view of the schedule - >>> # print(schedule.view()) - - ''' - - def __str__(self): - return "Changes a synchronous halo exchange into an asynchronous one." - - @property - def name(self): - ''' - :returns: the name of this transformation as a string. - :rtype: str - ''' - return "LFRicAsyncHaloExchangeTrans" - - def apply(self, node, options=None): - '''Transforms a synchronous halo exchange, represented by a - HaloExchange node, into an asynchronous halo exchange, - represented by HaloExchangeStart and HaloExchangeEnd nodes. - - :param node: a synchronous haloexchange node. - :type node: :py:obj:`psyclone.psygen.HaloExchange` - :param options: a dictionary with options for transformations. - :type options: Optional[Dict[str, Any]] - - ''' - self.validate(node, options) - - # add asynchronous start and end halo exchanges and initialise - # them using information from the existing synchronous halo - # exchange - # pylint: disable=protected-access - node.parent.addchild( - LFRicHaloExchangeStart( - node.field, check_dirty=node._check_dirty, - vector_index=node.vector_index, parent=node.parent), - index=node.position) - node.parent.addchild( - LFRicHaloExchangeEnd( - node.field, check_dirty=node._check_dirty, - vector_index=node.vector_index, parent=node.parent), - index=node.position) - - # remove the existing synchronous halo exchange - node.detach() - - def validate(self, node, options): - # pylint: disable=signature-differs - '''Internal method to check whether the node is valid for this - transformation. - - :param node: a synchronous Halo Exchange node - :type node: :py:obj:`psyclone.psygen.HaloExchange` - :param options: a dictionary with options for transformations. - :type options: Optional[Dict[str, Any]] - - :raises TransformationError: if the node argument is not a - HaloExchange (or subclass thereof) - - ''' - if not isinstance(node, psyGen.HaloExchange) or \ - isinstance(node, (LFRicHaloExchangeStart, LFRicHaloExchangeEnd)): - raise TransformationError( - f"Error in LFRicAsyncHaloExchange transformation. Supplied" - f" node must be a synchronous halo exchange but found " - f"'{type(node)}'.") - - class LFRicKernelConstTrans(Transformation): '''Modifies a kernel so that the number of dofs, number of layers and number of quadrature points are fixed in the kernel rather than @@ -2340,12 +2193,10 @@ def __new__(cls): "ACCParallelTrans", "ACCRoutineTrans", "ColourTrans", - "LFRicAsyncHaloExchangeTrans", "LFRicColourTrans", "LFRicKernelConstTrans", "LFRicOMPLoopTrans", "LFRicRedundantComputationTrans", - "LFRicOMPParallelLoopTrans", "KernelImportsToArguments", "MoveTrans", "OMPMasterTrans",