#2716 ensure Kern points to inlined PSyIR after transformation [skip ci]

src/psyclone/domain/common/transformations/kernel_module_inline_trans.py

@@ -385,10 +385,11 @@ def _rm_imported_symbol(name, table):
             actual_table.remove(csym)
 
     def apply(self, node, options=None):
-        ''' Bring the kernel subroutine into this Container.
+        ''' Bring the kernel/subroutine into this Container.
 
-        :param node: the kernel to module-inline.
-        :type node: :py:class:`psyclone.psyGen.CodedKern`
+        :param node: the Kernel or Call to module-inline.
+        :type node: :py:class:`psyclone.psyGen.CodedKern` |
+                    :py:class:`psyclone.psyir.nodes.Call`
         :param options: a dictionary with options for transformations.
         :type options: Optional[Dict[str, Any]]
 
@@ -412,10 +413,21 @@ def apply(self, node, options=None):
         # may already be in use, but the equality check below guarantees
         # that if it exists it is only valid when it references the exact same
         # implementation.
-        caller_name, codes_to_inline, interface_sym = (
+        # TODO - get rid of 'caller name' return value from
+        # _get_psyir_to_inline?
+        _, codes_to_inline, interface_sym = (
             KernelModuleInlineTrans._get_psyir_to_inline(node))
 
         updated_routines = self._prepare_code_to_inline(codes_to_inline)
+        # Update the Kernel to point to the updated PSyIR.
+        if isinstance(node, CodedKern):
+            # TODO - add setter for these properties to Kern?
+            # pylint: disable=protected-access
+            node._kern_schedule = updated_routines
+            if interface_sym:
+                node._interface_symbol = (
+                    updated_routines[0].scope.symbol_table.lookup(
+                        interface_sym.name))
 
         container = node.ancestor(Container)
         local_table = node.scope.symbol_table
@@ -475,7 +487,7 @@ def apply(self, node, options=None):
                             # is a detached copy.)
                             if routine != code_to_inline:
                                 raise TransformationError(
-                                    f"Cannot inline subroutine '{caller_name}'"
+                                    f"Cannot inline subroutine '{node.name}'"
                                     f" because another, different, subroutine "
                                     f"with the same name already exists and "
                                     f"versioning of module-inlined subroutines"

src/psyclone/domain/lfric/lfric_kern.py

@@ -53,7 +53,7 @@
 from psyclone.parse.algorithm import Arg, KernelCall
 from psyclone.psyGen import InvokeSchedule, CodedKern, args_filter
 from psyclone.psyir.frontend.fparser2 import Fparser2Reader
-from psyclone.psyir.nodes import (Loop, Literal, Reference,
+from psyclone.psyir.nodes import (Container, Loop, Literal, Reference,
                                   KernelSchedule)
 from psyclone.psyir.symbols import DataSymbol, ScalarType, ArrayType
 
@@ -653,72 +653,68 @@ class creates the PSyIR schedule on first invocation which is
 
         :raises GenerationError: if no subroutine matching this kernel can
             be found in the parse tree of the associated source code.
+
         '''
         if self._kern_schedule:
             return self._interface_symbol, self._kern_schedule
 
-        # Get the PSyIR Kernel Schedule(s)
-        routines = Fparser2Reader().get_routine_schedules(self.name, self.ast)
-        for routine in routines:
-            # If one of the symbols is not declared in a routine then
-            # this is only picked up when writing out the routine
-            # (raising a VisitorError), so we check here so that
-            # invalid code is not inlined. We use debug_string() to
-            # minimise the overhead.
-
-            # TODO #2271 could potentially avoid the need for
-            # debug_string() within. Sergi suggests that we may be
-            # missing the traversal of the declaration init
-            # expressions and that might solve the problem. I'm not so
-            # sure as we are talking about unknown symbols that will
-            # only be resolved in the back-end (or not). If I am right
-            # then one option would be to use the FortranWriter, but
-            # that would be bigger overhead, or perhaps just the
-            # declarations part of FortranWriter if that is possible.
-            # Also see TODO issue #2336 which captures the specific
-            # problem in LFRic that this fixes.
-            routine.debug_string()
-
-#       # The kernel name corresponds to an interface block. Find which
-#       # of the routines matches the precision of the arguments.
-#       for routine in routines:
-#           try:
-#               # The validity check for the kernel arguments should raise
-#               # an exception if the precisions don't match but currently
-#               # does not!
-#               self.validate_kernel_code_args(routine.symbol_table)
-#               sched = routine
-#               break
-#            except GenerationError:
-#               pass
-#        else:
-#            raise GenerationError(
-#                f"Failed to find a kernel implementation with an interface"
-#                f" that matches the invoke of '{self.name}'. (Tried "
-#                f"routines {[item.name for item in routines]}.)")
+        # Check for a local implementation of this kernel first.
+        container = self.ancestor(Container)
+        if container:
+            names = container.resolve_routine(self.name)
+            routines = []
+            for name in names:
+                rt_psyir = container.find_routine_psyir(name,
+                                                        allow_private=True)
+                routines.append(rt_psyir)
+
+        # Otherwise, get the PSyIR Kernel Schedule(s) from the original
+        # parse tree.
+        if not routines:
+            routines = Fparser2Reader().get_routine_schedules(self.name,
+                                                              self.ast)
+            new_schedules = []
+            for routine in routines[:]:
+                # If one of the symbols is not declared in a routine then
+                # this is only picked up when writing out the routine
+                # (raising a VisitorError), so we check here so that
+                # invalid code is not inlined. We use debug_string() to
+                # minimise the overhead.
+
+                # TODO #2271 could potentially avoid the need for
+                # debug_string() within. Sergi suggests that we may be
+                # missing the traversal of the declaration init
+                # expressions and that might solve the problem. I'm not so
+                # sure as we are talking about unknown symbols that will
+                # only be resolved in the back-end (or not). If I am right
+                # then one option would be to use the FortranWriter, but
+                # that would be bigger overhead, or perhaps just the
+                # declarations part of FortranWriter if that is possible.
+                # Also see TODO issue #2336 which captures the specific
+                # problem in LFRic that this fixes.
+                routine.debug_string()
+
+                # TODO #935 - replace the PSyIR argument data symbols with
+                # LFRic data symbols. For the moment we just return the
+                # unmodified PSyIR schedule but this should use
+                # RaisePSyIR2LFRicKernTrans once KernelInterface is fully
+                # functional (#928).
+                ksched = KernelSchedule(routine.symbol,
+                                        symbol_table=routine.symbol_table.detach())
+                for child in routine.pop_all_children():
+                    ksched.addchild(child)
+                routine.replace_with(ksched)
+                new_schedules.append(ksched)
+            routines = new_schedules
 
         if len(routines) > 1:
             table = routines[0].scope.symbol_table
             sym = table.lookup(self.name)
         else:
             sym = None
 
-        new_schedules = []
-        for sched in routines:
-            # TODO #935 - replace the PSyIR argument data symbols with
-            # LFRic data symbols. For the moment we just return the
-            # unmodified PSyIR schedule but this should use
-            # RaisePSyIR2LFRicKernTrans once KernelInterface is fully
-            # functional (#928).
-            ksched = KernelSchedule(sched.symbol,
-                                    symbol_table=sched.symbol_table.detach())
-            for child in sched.pop_all_children():
-                ksched.addchild(child)
-            sched.replace_with(ksched)
-            new_schedules.append(ksched)
-
         self._interface_symbol = sym
-        self._kern_schedule = new_schedules
+        self._kern_schedule = routines
 
         return self._interface_symbol, self._kern_schedule
 

src/psyclone/psyir/nodes/call.py

@@ -581,12 +581,14 @@ def _location_txt(node):
 
         if isinstance(container, Container):
             routines = []
-            for name in container.resolve_routine(rsym.name):
+            all_names = container.resolve_routine(rsym.name)
+            for name in all_names:
                 psyir = container.find_routine_psyir(
                     name, allow_private=can_be_private)
                 if psyir:
                     routines.append(psyir)
-            if routines:
+            if len(routines) == len(all_names):
+                # We've resolved everything.
                 return routines
 
         raise SymbolError(

src/psyclone/tests/domain/common/transformations/kernel_module_inline_trans_test.py

@@ -852,12 +852,15 @@ def test_psyir_mod_inline(fortran_reader, fortran_writer, tmpdir,
     assert "use my_mod, only : my_interface, my_other_sub\n" in output
     # We can't test the compilation of this code because of the 'use my_mod.'
 
+    # Finally, inline the call to the interface. This should then remove all
+    # imports from 'my_mod'.
     intrans.apply(calls[2])
     routines = container.walk(Routine)
     output = fortran_writer(psyir)
     assert "use my_mod" not in output
     assert "subroutine my_other_sub" in output
     assert "interface my_interface" in output
+    assert Compile(tmpdir).string_compiles(output)
 
 
 def test_mod_inline_no_container(fortran_reader, fortran_writer, tmpdir,