(related to #2244) LFRic tiled colouring

examples/lfric/scripts/tiledcolouring_and_omp.py

@@ -0,0 +1,80 @@
+# -----------------------------------------------------------------------------
+# BSD 3-Clause License
+#
+# Copyright (c) 2017-2023, 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 and S. Siso, STFC Daresbury Lab
+# Modified: I. Kavcic, Met Office
+# Modified by J. Henrichs, Bureau of Meteorology
+
+
+''' File containing a PSyclone transformation script for the Dynamo0p3
+API to apply colouring and OpenMP generically. This can be applied via
+the -s option in the "psyclone" script. '''
+from __future__ import print_function, absolute_import
+from psyclone.transformations import Dynamo0p3ColourTrans, \
+    DynamoOMPParallelLoopTrans
+from psyclone.psyir.nodes import Loop
+from psyclone.domain.lfric import LFRicConstants
+
+
+def trans(psy):
+    ''' PSyclone transformation script for the dynamo0p3 api to apply
+    colouring and OpenMP generically.'''
+    ctrans = Dynamo0p3ColourTrans()
+    otrans = DynamoOMPParallelLoopTrans()
+    const = LFRicConstants()
+
+    # Loop over all of the Invokes in the PSy object
+    for invoke in psy.invokes.invoke_list:
+
+        print("Transforming invoke '"+invoke.name+"'...")
+        schedule = invoke.schedule
+
+        # Colour all of the loops over cells unless they are on
+        # discontinuous spaces
+        for child in schedule.children:
+            if isinstance(child, Loop) \
+               and child.field_space.orig_name \
+               not in const.VALID_DISCONTINUOUS_NAMES \
+               and child.iteration_space == "cell_column":
+                ctrans.apply(child, options={"tiling": True})
+        # Then apply OpenMP to each of the colour loop
+        for child in schedule.children:
+            if isinstance(child, Loop):
+                if child.loop_type == "colours":
+                    otrans.apply(child.loop_body[0])
+                else:
+                    otrans.apply(child)
+
+        print(schedule.view())
+
+    return psy

src/psyclone/domain/lfric/kern_call_arg_list.py

@@ -964,28 +964,47 @@ def cell_ref_name(self, var_accesses=None):
         if self._kern.is_coloured():
             colour_sym = self._symtab.find_or_create_integer_symbol(
                 "colour", tag="colours_loop_idx")
+
+            from psyclone.domain.lfric import LFRicLoop
+            loop_type = self._kern.ancestor(LFRicLoop)._loop_type
+
+            # If there is only one colourmap we need to specify the tag
+            # to make sure we get the right symbol.
             if self._kern.is_intergrid:
                 tag = None
+            elif loop_type == "tile":
+                tag = "tmap"
             else:
-                # If there is only one colourmap we need to specify the tag
-                # to make sure we get the right symbol.
                 tag = "cmap"
-            array_ref = self.get_array_reference(self._kern.colourmap,
-                                                 [Reference(colour_sym),
-                                                  Reference(cell_sym)],
-                                                 ScalarType.Intrinsic.INTEGER,
-                                                 tag=tag)
-            if var_accesses is not None:
-                var_accesses.add_access(Signature(colour_sym.name),
-                                        AccessType.READ, self._kern)
-                var_accesses.add_access(Signature(cell_sym.name),
-                                        AccessType.READ, self._kern)
-                var_accesses.add_access(Signature(array_ref.name),
-                                        AccessType.READ,
-                                        self._kern, ["colour", "cell"])
-
-            return (self._kern.colourmap + "(colour,cell)",
-                    array_ref)
+
+            if loop_type == "tile":
+                tile_sym = self._symtab.find_or_create_integer_symbol(
+                    "tile", tag="tile_loop_idx")
+                array_ref = self.get_array_reference(
+                    self._kern.tilecolourmap,
+                    [Reference(colour_sym), Reference(tile_sym),
+                     Reference(cell_sym)],
+                    ScalarType.Intrinsic.INTEGER,
+                    tag=tag)
+                return (self._kern.tilecolourmap + "(colour,tile,cell)",
+                        array_ref)
+            else:
+                array_ref = self.get_array_reference(
+                                self._kern.colourmap,
+                                [Reference(colour_sym), Reference(cell_sym)],
+                                ScalarType.Intrinsic.INTEGER,
+                                tag=tag)
+                if var_accesses is not None:
+                    var_accesses.add_access(Signature(colour_sym.name),
+                                            AccessType.READ, self._kern)
+                    var_accesses.add_access(Signature(cell_sym.name),
+                                            AccessType.READ, self._kern)
+                    var_accesses.add_access(Signature(array_ref.name),
+                                            AccessType.READ,
+                                            self._kern, ["colour", "cell"])
+
+                return (self._kern.colourmap + "(colour,cell)",
+                        array_ref)
 
         if var_accesses is not None:
             var_accesses.add_access(Signature("cell"), AccessType.READ,

src/psyclone/domain/lfric/lfric_constants.py

@@ -168,8 +168,13 @@ def __init__(self):
         # halo. It is useful to group these together as we often need to
         # determine whether an access to a field or other object includes
         # access to the halo, or not.
-        LFRicConstants.HALO_ACCESS_LOOP_BOUNDS = ["cell_halo", "dof_halo",
-                                                  "colour_halo"]
+        LFRicConstants.HALO_ACCESS_LOOP_BOUNDS = [
+            "cell_halo",
+            "dof_halo",
+            "colour_halo",
+            "last_halo_tile_per_colour",
+            "last_halo_cell_per_colour_and_tile",
+        ]
 
         LFRicConstants.VALID_LOOP_BOUNDS_NAMES = (
             ["start",     # the starting
@@ -187,6 +192,10 @@ def __init__(self):
                           # the current colour
              "ncolours",  # the number of colours in a
                           # coloured loop
+             "ntilecolours",     # the number of colours in a
+                                 # coloured tiled loop
+             "last_edge_tile_per_colour",
+             "last_edge_cell_per_coloured_tile",
              "ncells",    # the number of owned cells
              "ndofs",     # the number of owned dofs
              "nannexed"]  # the number of owned dofs
@@ -205,7 +214,7 @@ def __init__(self):
         # (in the horizontal plane). A "null" loop doesn't iterate over
         # anything but is required for the halo-exchange logic.
         LFRicConstants.VALID_LOOP_TYPES = ["dof", "colours", "colour", "",
-                                           "null"]
+                                           "null", "colourtiles", "tile"]
 
         # Valid LFRic iteration spaces for built-in kernels
         LFRicConstants.BUILTIN_ITERATION_SPACES = ["dof"]

src/psyclone/domain/lfric/lfric_kern.py

@@ -458,7 +458,8 @@
         :returns: name of the colourmap (Fortran array).
         :rtype: str
 
-        :raises InternalError: if this kernel is not coloured.
+        :raises InternalError: if this kernel is not coloured or the dictionary
+        of inter-grid kernels and colourmaps has not been constructed.
 
         '''
         if not self.is_coloured():
@@ -479,6 +480,37 @@
 
         return cmap
 
+    @property
+    def tilecolourmap(self):
+        '''
+        Getter for the name of the tilecolourmap associated with this
+        kernel call.
+
+        :returns: name of the colourmap (Fortran array).
+        :rtype: str
+
+        :raises InternalError: if this kernel is not coloured or the dictionary
+        of inter-grid kernels and colourmaps has not been constructed.
+
+        '''
+        if not self.is_coloured():
+            raise InternalError(f"Kernel '{self.name}' is not inside a "
+                                f"coloured loop.")
+        sched = self.ancestor(InvokeSchedule)
+        if self.is_intergrid:
+            tmap = self._intergrid_ref.tilecolourmap_symbol.name
+        else:
+            try:
+                tmap = sched.symbol_table.lookup_with_tag("tmap").name
+            except KeyError:
+                # We have to do this here as _init_colourmap (which calls this
+                # method) is only called at code-generation time.
+                tmap = sched.symbol_table.find_or_create_array(
+                    "tmap", 3, ScalarType.Intrinsic.INTEGER,
+                    tag="tmap").name
+
+        return tmap
+
     @property
     def last_cell_all_colours_symbol(self):
         '''
@@ -533,6 +565,29 @@
 
         return self.scope.symbol_table.lookup_with_tag("ncolour").name
 
+    @property
+    def ntilecolours_var(self):
+        '''
+        Getter for the name of the variable holding the number of colours
+        associated with this kernel call.
+
+        :return: name of the variable holding the number of colours
+        :rtype: Union[str, NoneType]
+
+        :raises InternalError: if this kernel is not coloured or the
+            colour-map information has not been initialised.
+        '''
+        if not self.is_coloured():
+            raise InternalError(f"Kernel '{self.name}' is not inside a "
+                                f"coloured loop.")
+        if self.is_intergrid:
+            ncols_sym = self._intergrid_ref.ntilecolours_var_symbol
+            if not ncols_sym:
+                return None
+            return ncols_sym.name
+
+        return self.scope.symbol_table.lookup_with_tag("ntilecolour").name
+
     @property
     def fs_descriptors(self):
         '''