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

examples/gocean/eg3/ocl_trans.py

@@ -37,10 +37,10 @@
     the first Invoke to use OpenCL. '''
 
 from psyclone.psyGen import InvokeSchedule
-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(psyir):
@@ -62,7 +62,10 @@ def trans(psyir):
         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

@@ -91,7 +91,10 @@ def trans(psyir):
 
     for kernel in psyir.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/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, Routine
 from psyclone.psyir.transformations import (
@@ -76,6 +77,7 @@ def trans(psyir):
     otrans = Dynamo0p3OMPLoopTrans()
     const = LFRicConstants()
     cpu_parallel = OMPParallelTrans()
+    intrans = KernelModuleInlineTrans()
 
     if OFFLOAD_DIRECTIVES == "omp":
         # Use OpenMP offloading
@@ -120,7 +122,8 @@ def trans(psyir):
         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
@@ -137,6 +140,12 @@ def trans(psyir):
             if loop.iteration_space.endswith("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

@@ -55,11 +55,12 @@ def trans(psyir):
     # Loop over all of the Kernels Calls
     for kernel in psyir.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 "
+            print(f"Code of '{kernel.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