LFRic: (Closes #2189) fix bugs with missing maps in enter data directive

src/psyclone/domain/lfric/kern_call_acc_arg_list.py

@@ -42,7 +42,9 @@
 first before any members.
 '''
 
+from psyclone import psyGen
 from psyclone.domain.lfric import KernCallArgList
+from psyclone.errors import InternalError
 
 
 class KernCallAccArgList(KernCallArgList):
@@ -55,6 +57,42 @@ class KernCallAccArgList(KernCallArgList):
     to keep them in.
 
     '''
+    def cell_map(self, var_accesses=None):
+        '''Add cell-map to the list of required arrays.
+
+        :param var_accesses: optional VariablesAccessInfo instance to store
+            the information about variable accesses.
+        :type var_accesses: Optional[
+            :py:class:`psyclone.core.VariablesAccessInfo`]
+
+        '''
+        cargs = psyGen.args_filter(self._kern.args, arg_meshes=["gh_coarse"])
+        if len(cargs) > 1:
+            raise InternalError(
+                f"An LFRic intergrid kernel should have only one coarse mesh "
+                f"but '{self._kern.name}' has {len(cargs)}")
+        carg = cargs[0]
+        # Add the cell map to our argument list
+        base_name = "cell_map_" + carg.name
+        self.append(base_name)
+        # We'll need the current cell to index into this cell map.
+        self.cell_position(var_accesses)
+
+    def cell_position(self, var_accesses=None):
+        '''Adds a cell argument to the argument list and if supplied stores
+        this access in var_accesses. Although normally just a scalar, the cell
+        argument may actually require a lookup from a colour map array. Either
+        way, this method adds the name of the variable to the argument list.
+
+        :param var_accesses: optional VariablesAccessInfo instance to store
+            the information about variable accesses.
+        :type var_accesses: Optional[
+            :py:class:`psyclone.core.VariablesAccessInfo`]
+
+        '''
+        _, ref = self.cell_ref_name(var_accesses)
+        self.append(ref.symbol.name)
+
     def field_vector(self, argvect, var_accesses=None):
         '''Add the field vector associated with the argument 'argvect' to the
         argument list. OpenACC requires the field and the
@@ -63,10 +101,10 @@ def field_vector(self, argvect, var_accesses=None):
 
         :param argvect: the field vector to add.
         :type argvect: :py:class:`psyclone.dynamo0p3.DynKernelArgument`
-        :param var_accesses: optional VariablesAccessInfo instance to store \
+        :param var_accesses: optional VariablesAccessInfo instance to store
             the information about variable accesses.
-        :type var_accesses: \
-            :py:class:`psyclone.core.VariablesAccessInfo`
+        :type var_accesses: Optional[
+            :py:class:`psyclone.core.VariablesAccessInfo`]
 
         '''
         # First provide the derived-type object
@@ -83,10 +121,10 @@ def field(self, arg, var_accesses=None):
 
         :param arg: the field to be added.
         :type arg: :py:class:`psyclone.dynamo0p3.DynKernelArgument`
-        :param var_accesses: optional VariablesAccessInfo instance to store \
+        :param var_accesses: optional VariablesAccessInfo instance to store
             the information about variable accesses.
-        :type var_accesses: \
-            :py:class:`psyclone.core.VariablesAccessInfo`
+        :type var_accesses: Optional[
+            :py:class:`psyclone.core.VariablesAccessInfo`]
 
         '''
         text1 = arg.proxy_name
@@ -100,13 +138,13 @@ def stencil(self, arg, var_accesses=None):
         to the argument list. OpenACC requires the full dofmap to be
         specified. If supplied it also stores this access in var_accesses.
 
-        :param arg: the meta-data description of the kernel \
-            argument with which the stencil is associated.
+        :param arg: the meta-data description of the kernel argument with
+            which the stencil is associated.
         :type arg: :py:class:`psyclone.dynamo0p3.DynKernelArgument`
-        :param var_accesses: optional VariablesAccessInfo instance to store \
+        :param var_accesses: optional VariablesAccessInfo instance to store
             the information about variable accesses.
-        :type var_accesses: \
-            :py:class:`psyclone.core.VariablesAccessInfo`
+        :type var_accesses: Optional[
+            :py:class:`psyclone.core.VariablesAccessInfo`]
 
         '''
         # Import here to avoid circular dependency
@@ -122,13 +160,13 @@ def stencil_2d(self, arg, var_accesses=None):
         specified. If supplied it also stores this access in var_accesses.This
         method passes through to the stencil method.
 
-        :param arg: the meta-data description of the kernel \
+        :param arg: the meta-data description of the kernel
             argument with which the stencil is associated.
         :type arg: :py:class:`psyclone.dynamo0p3.DynKernelArgument`
-        :param var_accesses: optional VariablesAccessInfo instance to store \
+        :param var_accesses: optional VariablesAccessInfo instance to store
             the information about variable accesses.
-        :type var_accesses: \
-            :py:class:`psyclone.core.VariablesAccessInfo`
+        :type var_accesses: Optional[
+            :py:class:`psyclone.core.VariablesAccessInfo`]
 
         '''
         self.stencil(arg, var_accesses)
@@ -140,10 +178,10 @@ def stencil_unknown_extent(self, arg, var_accesses=None):
 
         :param arg: the kernel argument with which the stencil is associated.
         :type arg: :py:class:`psyclone.dynamo0p3.DynKernelArgument`
-        :param var_accesses: optional VariablesAccessInfo instance to store \
+        :param var_accesses: optional VariablesAccessInfo instance to store
             the information about variable accesses.
-        :type var_accesses: \
-            :py:class:`psyclone.core.VariablesAccessInfo`
+        :type var_accesses: Optional[
+            :py:class:`psyclone.core.VariablesAccessInfo`]
 
         '''
         # The extent is not specified in the metadata so pass the value in
@@ -162,10 +200,10 @@ def stencil_2d_unknown_extent(self, arg, var_accesses=None):
 
         :param arg: the kernel argument with which the stencil is associated.
         :type arg: :py:class:`psyclone.dynamo0p3.DynKernelArgument`
-        :param var_accesses: optional VariablesAccessInfo instance to store \
+        :param var_accesses: optional VariablesAccessInfo instance to store
             the information about variable accesses.
-        :type var_accesses: \
-            :py:class:`psyclone.core.VariablesAccessInfo`
+        :type var_accesses: Optional[
+            :py:class:`psyclone.core.VariablesAccessInfo`]
 
         '''
         self.stencil_unknown_extent(arg, var_accesses)
@@ -178,10 +216,10 @@ def operator(self, arg, var_accesses=None):
 
         :param arg: the meta-data description of the operator.
         :type arg: :py:class:`psyclone.dynamo0p3.DynKernelArgument`
-        :param var_accesses: optional VariablesAccessInfo instance to store \
+        :param var_accesses: optional VariablesAccessInfo instance to store
             the information about variable accesses.
-        :type var_accesses: \
-            :py:class:`psyclone.core.VariablesAccessInfo`
+        :type var_accesses: Optional[
+            :py:class:`psyclone.core.VariablesAccessInfo`]
 
         '''
         # In case of OpenACC we do not want to transfer the same
@@ -197,13 +235,13 @@ def fs_compulsory_field(self, function_space, var_accesses=None):
         to be specified. If supplied it also stores this access in
         var_accesses.
 
-        :param function_space: the function space for which the compulsory \
+        :param function_space: the function space for which the compulsory
             arguments are added.
         :type function_space: :py:class:`psyclone.domain.lfric.FunctionSpace`
-        :param var_accesses: optional VariablesAccessInfo instance to store \
+        :param var_accesses: optional VariablesAccessInfo instance to store
             the information about variable accesses.
-        :type var_accesses: \
-            :py:class:`psyclone.core.VariablesAccessInfo`
+        :type var_accesses: Optional[
+            :py:class:`psyclone.core.VariablesAccessInfo`]
 
         '''
         if self._kern.iterates_over != "cell_column":
@@ -213,17 +251,42 @@ def fs_compulsory_field(self, function_space, var_accesses=None):
         # needs the whole field, so we cannot call the base class.
         self.append(function_space.map_name, var_accesses)
 
+    def fs_intergrid(self, function_space, var_accesses=None):
+        '''Add arrays that need to be uploaded for inter-grid kernels.
+        These arrays contain the mapping between fine and coarse meshes.
+
+        :param function_space: the function space associated with the mesh.
+        :type function_space: :py:class:`psyclone.domain.lfric.FunctionSpace`
+        :param var_accesses: optional VariablesAccessInfo instance to store
+            the information about variable accesses.
+        :type var_accesses: Optional[
+            :py:class:`psyclone.core.VariablesAccessInfo`]
+
+        '''
+        # Is this FS associated with the coarse or fine mesh? (All fields
+        # on a given mesh must be on the same FS.)
+        arg = self._kern.arguments.get_arg_on_space(function_space)
+        if arg.mesh == "gh_fine":
+            # For the fine mesh, we need the *whole* dofmap
+            map_name = function_space.map_name
+            self.append(map_name, var_accesses)
+        else:
+            # For the coarse mesh we only need undf and the dofmap for
+            # the current column
+            self.fs_compulsory_field(function_space,
+                                     var_accesses=var_accesses)
+
     def scalar(self, scalar_arg, var_accesses=None):
         '''
         Override the default implementation as there's no need to specify
         scalars for an OpenACC data region.
 
         :param scalar_arg: the kernel argument.
         :type scalar_arg: :py:class:`psyclone.dynamo0p3.DynKernelArgument`
-        :param var_accesses: optional VariablesAccessInfo instance that \
+        :param var_accesses: optional VariablesAccessInfo instance that
             stores information about variable accesses.
-        :type var_accesses: \
-            :py:class:`psyclone.core.VariablesAccessInfo`
+        :type var_accesses: Optional[
+            :py:class:`psyclone.core.VariablesAccessInfo`]
 
         '''
 

src/psyclone/domain/lfric/kern_call_arg_list.py

@@ -993,10 +993,17 @@ 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")
-            array_ref = self.get_array_reference("cmap",
+            if self._kern.is_intergrid:
+                tag = None
+            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)
+                                                 ScalarType.Intrinsic.INTEGER,
+                                                 tag=tag)
             if var_accesses is not None:
                 var_accesses.add_access(Signature(colour_sym.name),
                                         AccessType.READ, self._kern)

src/psyclone/dynamo0p3.py

@@ -1860,7 +1860,8 @@ def kern_args(self, stub=False, var_accesses=None,
                     cell_name = "cell"
                     if self._kernel.is_coloured():
                         colour_name = "colour"
-                        cmap_name = "cmap"
+                        cmap_name = self._symbol_table.find_or_create_tag(
+                            "cmap", root_name="cmap").name
                         adj_face += (f"(:,{cmap_name}({colour_name},"
                                      f"{cell_name}))")
                     else:
@@ -3720,12 +3721,13 @@ def _colourmap_init(self):
         '''
         Sets-up information on any required colourmaps. This cannot be done
         in the constructor since colouring is applied by Transformations
-        and happens after the Schedule has already been constructed.
+        and happens after the Schedule has already been constructed. Therefore,
+        this method is called at code-generation time.
 
         '''
         # pylint: disable=too-many-locals
         const = LFRicConstants()
-        have_non_intergrid = False
+        non_intergrid_kern = None
         sym_tab = self._schedule.symbol_table
 
         for call in [call for call in self._schedule.coded_kernels() if
@@ -3740,7 +3742,7 @@ def _colourmap_init(self):
                 self._needs_colourmap = True
 
             if not call.is_intergrid:
-                have_non_intergrid = True
+                non_intergrid_kern = call
                 continue
 
             # This is an inter-grid kernel so look-up the names of
@@ -3775,13 +3777,14 @@ def _colourmap_init(self):
             self._ig_kernels[id(call)].set_colour_info(
                 colour_map, ncolours, last_cell)
 
-        if have_non_intergrid and (self._needs_colourmap or
+        if non_intergrid_kern and (self._needs_colourmap or
                                    self._needs_colourmap_halo):
             # There aren't any inter-grid kernels but we do need colourmap
             # information and that means we'll need a mesh object
             self._add_mesh_symbols(["mesh"])
-            colour_map = sym_tab.find_or_create_array(
-                "cmap", 2, ScalarType.Intrinsic.INTEGER, tag="cmap").name
+            # This creates the colourmap information for this invoke if we
+            # don't already have one.
+            colour_map = non_intergrid_kern.colourmap
             # No. of colours
             ncolours = sym_tab.find_or_create_integer_symbol(
                 "ncolour", tag="ncolour").name
@@ -3881,8 +3884,8 @@ def declarations(self, parent):
             # There aren't any inter-grid kernels but we do need
             # colourmap information
             base_name = "cmap"
-            colour_map = \
-                self._schedule.symbol_table.find_or_create_tag(base_name).name
+            csym = self._schedule.symbol_table.lookup_with_tag("cmap")
+            colour_map = csym.name
             # No. of colours
             base_name = "ncolour"
             ncolours = \
@@ -7736,16 +7739,24 @@ def colourmap(self):
         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:
-            invoke = self.ancestor(InvokeSchedule).invoke
+            invoke = sched.invoke
             if id(self) not in invoke.meshes.intergrid_kernels:
                 raise InternalError(
                     f"Colourmap information for kernel '{self.name}' has "
                     f"not yet been initialised")
             cmap = invoke.meshes.intergrid_kernels[id(self)].\
                 colourmap_symbol.name
         else:
-            cmap = self.scope.symbol_table.lookup_with_tag("cmap").name
+            try:
+                cmap = sched.symbol_table.lookup_with_tag("cmap").name
+            except KeyError:
+                # We have to do this here as _init_colourmap (which calls this
+                # method) is only called at code-generation time.
+                cmap = sched.symbol_table.find_or_create_array(
+                    "cmap", 2, ScalarType.Intrinsic.INTEGER,
+                    tag="cmap").name
 
         return cmap