(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,10 +43,14 @@
 '''
 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.psyGen import CodedKern
+from psyclone.psyir.nodes import (
+    Call, Directive, IntrinsicCall, Loop, Routine, Schedule)
 from psyclone.psyir.transformations import (
-    ACCKernelsTrans, TransformationError, OMPTargetTrans)
+    ACCKernelsTrans, InlineTrans, Matmul2CodeTrans, OMPTargetTrans,
+    TransformationError)
 from psyclone.transformations import (
     Dynamo0p3ColourTrans, Dynamo0p3OMPLoopTrans,
     Dynamo0p3RedundantComputationTrans, OMPParallelTrans,
@@ -58,9 +62,54 @@
 INVOKE_EXCLUSIONS = [
 ]
 
+# We won't attempt to inline calls to routines with names that contain
+# these strings.
+INLINE_EXCLUSIONS = ["abort", "logging"]
+
 OFFLOAD_DIRECTIVES = os.getenv('LFRIC_OFFLOAD_DIRECTIVES', "none")
 
 
+def _inline_calls(kern):
+    '''
+    Recursively inline all calls within the supplied Kernel or Routine.
+
+    :param kern: the Kernel or Routine to inline any Calls into.
+
+    '''
+    mod_inline_trans = KernelModuleInlineTrans()
+    intrans = InlineTrans()
+    matrans = Matmul2CodeTrans()
+
+    if isinstance(kern, CodedKern):
+        _, scheds = kern.get_kernel_schedule()
+    else:
+        scheds = [kern]
+    for sched in scheds:
+        sched: Schedule
+        for call in sched.walk(Call):
+            call: Call
+            if isinstance(call, IntrinsicCall):
+                try:
+                    matrans.apply(call)
+                except TransformationError:
+                    pass
+                continue
+            if any(name in call.routine.name for name in INLINE_EXCLUSIONS):
+                continue
+            try:
+                for inner_call in call.get_callees():
+                    _inline_calls(inner_call)
+                mod_inline_trans.apply(call)
+                try:
+                    intrans.apply(call)
+                except TransformationError as err:
+                    print(f"Failed to inline call {call.debug_string()}:\n"
+                          f"{err}")
+            except (TransformationError, NotImplementedError) as err:
+                print(f"Failed to module-inline routine {call.routine.name}:\n"
+                      f"{err}")
+
+
 def trans(psyir):
     '''Applies PSyclone colouring and GPU offloading transformations. Any
     kernels that cannot be offloaded to GPU are parallelised using OpenMP
@@ -76,6 +125,7 @@ def trans(psyir):
     otrans = Dynamo0p3OMPLoopTrans()
     const = LFRicConstants()
     cpu_parallel = OMPParallelTrans()
+    mod_inline_trans = KernelModuleInlineTrans()
 
     if OFFLOAD_DIRECTIVES == "omp":
         # Use OpenMP offloading
@@ -120,7 +170,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 +188,20 @@ def trans(psyir):
             if loop.iteration_space.endswith("cell_column"):
                 if offload:
                     for kern in loop.kernels():
+                        try:
+                            mod_inline_trans.apply(kern)
+                            _inline_calls(kern)
+                            # At this point we would like to fully inline the
+                            # kernel but InlineTrans does not accept a
+                            # CodedKern. If we lower this kernel first then we
+                            # get errors later (at code-generation time).
+                            # Hopefully this will be resolved when we move
+                            # LFRic to use the PSyIR backend for code
+                            # generation.
+                        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