(Closes #2716) Add support for module-inlining calls to polymorphic kernels/routines

examples/gocean/eg3/ocl_trans.py

@@ -36,10 +36,10 @@
 ''' Module providing a transformation script that converts the Schedule of
     the first Invoke to use OpenCL. '''
 
-from psyclone.psyir.transformations import \
-    FoldConditionalReturnExpressionsTrans
-from psyclone.domain.gocean.transformations import GOOpenCLTrans, \
-    GOMoveIterationBoundariesInsideKernelTrans
+from psyclone.psyir.transformations import (
+    FoldConditionalReturnExpressionsTrans)
+from psyclone.domain.gocean.transformations import (
+    GOOpenCLTrans, GOMoveIterationBoundariesInsideKernelTrans)
 
 
 def trans(psy):
@@ -68,7 +68,10 @@ def trans(psy):
         move_boundaries_trans.apply(kern)
         # Change the syntax to remove the return statements introduced by the
         # previous transformation
-        fold_trans.apply(kern.get_kernel_schedule())
+        _, kschedules = kern.get_kernel_schedule()
+        # NOTE: we assume the kernel is not polymorphic and thus there is
+        # only one schedule associated with it.
+        fold_trans.apply(kschedules[0])
         # Specify the OpenCL queue and workgroup size of the kernel
         # In this case we dispatch each kernel in a different queue to check
         # that the output code has the necessary barriers to guarantee the

examples/lfric/eg15/matvec_opt.py

@@ -75,7 +75,7 @@
 
 
 def trans(psy):
-    '''PSyclone transformation script for the Dynamo0.3 API to optimise
+    '''PSyclone transformation script for the LFRic API to optimise
     the matvec kernel for many-core CPUs. For the moment simply find
     the first matvec kernel in the example, transform the matmul
     intrinsic to equivalant inline code and then print out its PSyIR
@@ -90,7 +90,10 @@ def trans(psy):
         schedule = invoke.schedule
         for kernel in schedule.coded_kernels():
             if kernel.name.lower() == "matrix_vector_kernel_code":
-                kernel_schedule = kernel.get_kernel_schedule()
+                _, kernel_schedules = kernel.get_kernel_schedule()
+                # For simplicity, ASSUME that the kernel is not polymorphic and
+                # thus only has one schedule.
+                kernel_schedule = kernel_schedules[0]
                 # Replace matmul with inline code
                 for icall in kernel_schedule.walk(IntrinsicCall):
                     if icall.intrinsic is IntrinsicCall.Intrinsic.MATMUL:

examples/lfric/eg2/module_inline_trans.py

@@ -46,13 +46,10 @@ def trans(psy):
     PSyclone transformation routine. This is an example which module-inlines
     the kernel used in the second 'invoke' in the supplied PSy object.
 
-    :param psy: the PSy object that PSyclone has constructed for the \
+    :param psy: the PSy object that PSyclone has constructed for the
                 'invoke'(s) found in the Algorithm file.
     :type psy: :py:class:`psyclone.dynamo0p3.DynamoPSy`
 
-    :returns: the transformed PSy object.
-    :rtype: :py:class:`psyclone.dynamo0p3.DynamoPSy`
-
     '''
     invokes = psy.invokes
     print(psy.invokes.names)
@@ -61,15 +58,7 @@ def trans(psy):
     print(schedule.view())
     # Find the kernel we want to inline.
     kern = schedule.walk(Kern)[0]
-    # Setting module inline directly.
-    kern.module_inline = True
-    print(schedule.view())
-    # Unsetting module inline via a transformation.
     inline_trans = KernelModuleInlineTrans()
-    inline_trans.apply(kern, {"inline": False})
-    print(schedule.view())
-    # Setting module inline via a transformation.
+    # Apply the inlining transformation.
     inline_trans.apply(kern)
     print(schedule.view())
-
-    return psy

examples/lfric/scripts/gpu_offloading.py

@@ -43,6 +43,7 @@
 '''
 import os
 import sys
+from psyclone.domain.common.transformations import KernelModuleInlineTrans
 from psyclone.domain.lfric import LFRicConstants
 from psyclone.psyir.nodes import Directive, Loop
 from psyclone.psyir.transformations import (
@@ -73,6 +74,7 @@ def trans(psy):
     otrans = Dynamo0p3OMPLoopTrans()
     const = LFRicConstants()
     cpu_parallel = OMPParallelTrans()
+    intrans = KernelModuleInlineTrans()
 
     if OFFLOAD_DIRECTIVES == "omp":
         # Use OpenMP offloading
@@ -119,7 +121,8 @@ def trans(psy):
         else:
             offload = True
 
-        # Keep a record of any kernels we fail to offload
+        # Keep a record of any kernels we fail to offload.
+        failed_inline = set()
         failed_to_offload = set()
 
         # Colour loops over cells unless they are on discontinuous spaces
@@ -136,6 +139,12 @@ def trans(psy):
             if loop.iteration_space == "cell_column":
                 if offload:
                     for kern in loop.kernels():
+                        try:
+                            intrans.apply(kern)
+                        except TransformationError as err:
+                            failed_inline.add(kern.name.lower())
+                            print(f"Failed to module-inline kernel "
+                                  f"'{kern.name}' due to:\n{err.value}")
                         try:
                             gpu_annotation_trans.apply(kern)
                         except TransformationError as err:

examples/lfric/scripts/kernel_print.py

@@ -54,11 +54,12 @@ def trans(psy):
         # Loop over all of the Kernels in this Schedule.
         for kernel in schedule.coded_kernels():
             try:
-                kernel_schedule = kernel.get_kernel_schedule()
-                if kernel_schedule not in already_printed:
-                    kern = fortran_writer(kernel_schedule)
-                    print(kern)
-                    already_printed.append(kernel_schedule)
+                _, kernel_schedules = kernel.get_kernel_schedule()
+                for ksched in kernel_schedules:
+                    if ksched not in already_printed:
+                        kern = fortran_writer(ksched)
+                        print(kern)
+                    already_printed.append(ksched)
             except Exception as err:  # pylint: disable=broad-except
                 print(f"Code of '{kernel.name}' in '{invoke.name}' "
                       f"cannot be printed because:\n{err}")

examples/xdsl/backend/xdsl.py

@@ -439,7 +439,9 @@ def checkIfStringIsType(self, string, typ):
 
     def nemokern_node(self, node):
         exec_statements = []
-        schedule = node.get_kernel_schedule()
+        _, schedules = node.get_kernel_schedule()
+        # IGNORE polymorphic routines.
+        schedule = schedules[0]
         for child in schedule.children:
             exec_statements.append(self._visit(child))
         return exec_statements