Merge pull request #415 from ecmwf-ifs/naml-transformations-parallel

Transformations: "Parallel" sub-module with driver-level parallelilsation utilities

loki/ir/pragma_utils.py

@@ -532,7 +532,7 @@ def visit_tuple(self, o, **kwargs):
             if start in o:
                 # If a pair does not live in the same tuple we have a problem.
                 if stop not in o:
-                    warning('[Loki::IR] Cannot find matching end for pragma {start} at same IR level!')
+                    warning(f'[Loki::IR] Cannot find matching end for pragma {start} at same IR level!')
                     continue
 
                 # Create the PragmaRegion node and replace in tuple

loki/transformations/__init__.py

@@ -36,3 +36,4 @@
 from loki.transformations.split_read_write import * # noqa
 from loki.transformations.loop_blocking import * # noqa
 from loki.transformations.routine_signatures import * # noqa
+from loki.transformations.parallel import * # noqa

loki/transformations/parallel/__init__.py

@@ -0,0 +1,15 @@
+# (C) Copyright 2018- ECMWF.
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+"""
+Sub-package with utilities to remove, generate and manipulate parallel
+regions.
+"""
+
+from loki.transformations.parallel.block_loop import * # noqa
+from loki.transformations.parallel.field_api import * # noqa
+from loki.transformations.parallel.openmp_region import * # noqa

loki/transformations/parallel/block_loop.py

@@ -0,0 +1,137 @@
+# (C) Copyright 2018- ECMWF.
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+"""
+Transformation utilities to remove and generate parallel block loops.
+"""
+
+from loki.expression import symbols as sym
+from loki.ir import (
+    nodes as ir, FindNodes, Transformer, pragma_regions_attached,
+    is_loki_pragma
+)
+from loki.tools import as_tuple
+from loki.types import BasicType, SymbolAttributes
+
+
+__all__ = ['remove_block_loops', 'add_block_loops']
+
+
+def remove_block_loops(routine, dimension):
+    """
+    Remove any outer block :any:`Loop` from a given :any:`Subroutine.
+
+    The loops are identified according to a given :any:`Dimension`
+    object, and will remove auxiliary assignments of index and bound
+    variables, as commonly used in IFS-style block loops.
+
+    Parameters
+    ----------
+    routine: :any:`Subroutine`
+        Subroutine from which to remove block loops
+    dimension : :any:`Dimension`
+        The dimension object describing loop variables
+    """
+    idx = dimension.index
+    variables = as_tuple(dimension.indices)
+    variables += as_tuple(dimension.lower)
+    variables += as_tuple(dimension.upper)
+
+    class RemoveBlockLoopTransformer(Transformer):
+        """
+        :any:`Transformer` to remove driver-level block loops.
+        """
+
+        def visit_Loop(self, loop, **kwargs):  # pylint: disable=unused-argument
+            body = self.visit(loop.body, **kwargs)
+
+            if not loop.variable == idx:
+                return loop._rebuild(body=body)
+
+            to_remove = tuple(
+                a for a in FindNodes(ir.Assignment).visit(body)
+                if a.lhs in variables
+            )
+            return tuple(n for n in body if n not in to_remove)
+
+    routine.body = RemoveBlockLoopTransformer().visit(routine.body)
+
+
+def add_block_loops(routine, dimension, default_type=None):
+    """
+    Insert IFS-style (NPROMA) driver block-loops in ``!$loki
+    parallel`` regions.
+
+    The provided :any:`Dimension` object describes the variables to
+    used when generating the loop and default assignments. It
+    encapsulates IFS-specific convention, where a strided loop over
+    points, defined by ``dimension.index``, ``dimension.bounds`` and
+    ``dimension.step`` is created, alongside assignments that define
+    the corresponding block index and upper bound, defined by
+    ``dimension.indices[1]`` and ``dimension.upper[1]`` respectively.
+
+    Parameters
+    ----------
+    routine : :any:`Subroutine`
+        The routine in which to add block loops.
+    dimension : :any:`Dimension`
+        The dimension object describing the block loop variables.
+    default_type : :any:`SymbolAttributes`, optional
+        Default type to use when creating variables; defaults to
+        ``integer(kind=JPIM)``.
+    """
+
+    _default = SymbolAttributes(BasicType.INTEGER, kind='JPIM')
+    dtype = default_type if default_type else _default
+
+    lidx = routine.parse_expr(dimension.index)
+    bidx = routine.parse_expr(dimension.indices[1])
+    bupper = routine.parse_expr(dimension.upper[1])
+
+    # Ensure that local integer variables are declared
+    for v in (lidx, bupper, bidx):
+        if not v in routine.variable_map:
+            routine.variables += (v.clone(type=dtype),)
+
+    def _create_block_loop(body, scope):
+        """
+        Generate block loop object, including indexing preamble
+        """
+
+        bsize = scope.parse_expr(dimension.step)
+        lupper = scope.parse_expr(dimension.upper[0])
+        lrange = sym.LoopRange((sym.Literal(1), lupper, bsize))
+
+        expr_tail = scope.parse_expr(f'{lupper}-{lidx}+1')
+        expr_max = sym.InlineCall(
+            function=sym.ProcedureSymbol('MIN', scope=scope), parameters=(bsize, expr_tail)
+        )
+        preamble = (ir.Assignment(lhs=bupper, rhs=expr_max),)
+        preamble += (ir.Assignment(
+            lhs=bidx, rhs=scope.parse_expr(f'({lidx}-1)/{bsize}+1')
+        ),)
+
+        return ir.Loop(variable=lidx, bounds=lrange, body=preamble + body)
+
+    class InsertBlockLoopTransformer(Transformer):
+
+        def visit_PragmaRegion(self, region, **kwargs):
+            """
+            (Re-)insert driver-level block loops into marked parallel region.
+            """
+            if not is_loki_pragma(region.pragma, starts_with='parallel'):
+                return region
+
+            scope = kwargs.get('scope')
+
+            loop = _create_block_loop(body=region.body, scope=scope)
+
+            region._update(body=(ir.Comment(''), loop))
+            return region
+
+    with pragma_regions_attached(routine):
+        routine.body = InsertBlockLoopTransformer().visit(routine.body, scope=routine)

loki/transformations/parallel/field_api.py

@@ -0,0 +1,152 @@
+# (C) Copyright 2018- ECMWF.
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+"""
+Transformation utilities to manage and inject FIELD-API boilerplate code.
+"""
+
+from loki.expression import symbols as sym
+from loki.ir import (
+    nodes as ir, FindNodes, FindVariables, Transformer
+)
+from loki.logging import warning
+from loki.tools import as_tuple
+
+
+__all__ = [
+    'remove_field_api_view_updates', 'add_field_api_view_updates'
+]
+
+
+def remove_field_api_view_updates(routine, field_group_types, dim_object=None):
+    """
+    Remove FIELD API boilerplate calls for view updates of derived types.
+
+    This utility is intended to remove the IFS-specific group type
+    objects that provide block-scope view pointers to deep kernel
+    trees. It will remove all calls to ``UPDATE_VIEW`` on derive-type
+    objects with the respective types.
+
+    Parameters
+    ----------
+    routine : :any:`Subroutine`
+        The routine from which to remove FIELD API update calls
+    field_group_types : tuple of str
+        List of names of the derived types of "field group" objects to remove
+    dim_object : str, optional
+        Optional name of the "dimension" object; if provided it will remove the
+        call to ``<dim>%UPDATE(...)`` accordingly.
+    """
+    field_group_types = as_tuple(str(fgt).lower() for fgt in field_group_types)
+
+    class RemoveFieldAPITransformer(Transformer):
+
+        def visit_CallStatement(self, call, **kwargs):  # pylint: disable=unused-argument
+
+            if '%update_view' in str(call.name).lower():
+                if not str(call.name.parent.type.dtype).lower() in field_group_types:
+                    warning(f'[Loki::ControlFlow] Removing {call.name} call, but not in field group types!')
+
+                return None
+
+            if dim_object and f'{dim_object}%update'.lower() in str(call.name).lower():
+                return None
+
+            return call
+
+        def visit_Assignment(self, assign, **kwargs):  # pylint: disable=unused-argument
+            if str(assign.lhs.type.dtype).lower() in field_group_types:
+                warning(f'[Loki::ControlFlow] Found LHS field group assign: {assign}')
+            return assign
+
+        def visit_Loop(self, loop, **kwargs):
+            loop = self.visit_Node(loop, **kwargs)
+            return loop if loop.body else None
+
+        def visit_Conditional(self, cond, **kwargs):
+            cond = super().visit_Node(cond, **kwargs)
+            return cond if cond.body else None
+
+    routine.body = RemoveFieldAPITransformer().visit(routine.body)
+
+
+def add_field_api_view_updates(routine, dimension, field_group_types, dim_object=None):
+    """
+    Adds FIELD API boilerplate calls for view updates.
+
+    The provided :any:`Dimension` object describes the local loop variables to
+    pass to the respective update calls. In particular, ``dimension.indices[1]``
+    is used to denote the block loop index that is passed to ``UPDATE_VIEW()``
+    calls on field group object. The list of type names ``field_group_types``
+    is used to identify for which objcets the view update calls get added.
+
+    Parameters
+    ----------
+    routine : :any:`Subroutine`
+        The routine from which to remove FIELD API update calls
+    dimension : :any:`Dimension`
+        The dimension object describing the block loop variables.
+    field_group_types : tuple of str
+        List of names of the derived types of "field group" objects to remove
+    dim_object : str, optional
+        Optional name of the "dimension" object; if provided it will remove the
+        call to ``<dim>%UPDATE(...)`` accordingly.
+    """
+
+    def _create_dim_update(scope, dim_object):
+        index = scope.parse_expr(dimension.index)
+        upper = scope.parse_expr(dimension.upper[1])
+        bindex = scope.parse_expr(dimension.indices[1])
+        idims = scope.get_symbol(dim_object)
+        csym = sym.ProcedureSymbol(name='UPDATE', parent=idims, scope=idims.scope)
+        return ir.CallStatement(name=csym, arguments=(bindex, upper, index), kwarguments=())
+
+    def _create_view_updates(section, scope):
+        bindex = scope.parse_expr(dimension.indices[1])
+
+        fgroup_vars = sorted(tuple(
+            v for v in FindVariables(unique=True).visit(section)
+            if str(v.type.dtype) in field_group_types
+        ), key=str)
+        calls = ()
+        for fgvar in fgroup_vars:
+            fgsym = scope.get_symbol(fgvar.name)
+            csym = sym.ProcedureSymbol(name='UPDATE_VIEW', parent=fgsym, scope=fgsym.scope)
+            calls += (ir.CallStatement(name=csym, arguments=(bindex,), kwarguments=()),)
+
+        return calls
+
+    class InsertFieldAPIViewsTransformer(Transformer):
+        """ Injects FIELD-API view updates into block loops """
+
+        def visit_Loop(self, loop, **kwargs):  # pylint: disable=unused-argument
+            if not loop.variable == 'JKGLO':
+                return loop
+
+            scope = kwargs.get('scope')
+
+            # Find the loop-setup assignments
+            _loop_symbols = dimension.indices
+            _loop_symbols += as_tuple(dimension.lower) + as_tuple(dimension.upper)
+            loop_setup = tuple(
+                a for a in FindNodes(ir.Assignment).visit(loop.body)
+                if a.lhs in _loop_symbols
+            )
+            idx = max(loop.body.index(a) for a in loop_setup) + 1
+
+            # Prepend FIELD API boilerplate
+            preamble = (
+                ir.Comment(''), ir.Comment('! Set up thread-local view pointers')
+            )
+            if dim_object:
+                preamble += (_create_dim_update(scope, dim_object=dim_object),)
+            preamble += _create_view_updates(loop.body, scope)
+
+            loop._update(body=loop.body[:idx] + preamble + loop.body[idx:])
+            return loop
+
+    routine.body = InsertFieldAPIViewsTransformer().visit(routine.body, scope=routine)