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

changelog

@@ -1,3 +1,6 @@
+   43) PR #2732 for #2716. Add support for module-inlining polymorphic kernels
+   and rename get_kernel_schedule to get_callees (to match the Call method).
+
    42) PR #3034 for #2837. Updates the supported versions of Python used in
    the test suite to 3.10 and 3.13.
 

doc/developer_guide/transformations.rst

@@ -51,13 +51,13 @@ Kernel Transformations
 PSyclone is able to perform kernel transformations by obtaining the PSyIR
 representation of the kernel with:
 
-.. automethod:: psyclone.psyGen.CodedKern.get_kernel_schedule
+.. automethod:: psyclone.psyGen.CodedKern.get_callees
     :no-index:
 
-The result of `psyclone.psyGen.Kern.get_kernel_schedule` is a
-`psyclone.psyir.nodes.KernelSchedule` which is a specialisation of the
-`Routine` class with the `is_program` and `return_type` properties set to
-`False` and `None`, respectively.
+The result of `psyclone.psyGen.Kern.get_callees` is a list of
+`psyclone.psyir.nodes.KernelSchedule` objects. `KernelSchedule` is a
+specialisation of the `Routine` class with the `is_program` and `return_type`
+properties set to False` and `None`, respectively.
 
 In addition to modifying the kernel PSyIR with the desired transformations,
 the `modified` flag of the `CodedKern` node has to be set. This will let

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_callees()
+        # 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_callees()
+            # 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/eg19/Makefile

@@ -59,8 +59,11 @@ $(EXEC): $(LFRIC_LIB) $(OBJ)
 
 compile: transform $(EXEC)
 
+# Runs PSyclone to do the code generation. Also demonstrates the use of the
+# example 'kernel_print' transformation which prints the Fortran of each
+# kernel found.
 transform:
-	${PSYCLONE} -api lfric algorithm.x90 -opsy mixed_precision_psy.f90 -oalg alg.f90
+	${PSYCLONE} -api lfric algorithm.x90 -s ../scripts/kernel_print.py -opsy mixed_precision_psy.f90 -oalg alg.f90
 
 alg.f90 mixed_precision_psy.f90: transform
 alg.o: mixed_precision_psy.o

examples/lfric/scripts/gpu_offloading.py

@@ -45,9 +45,10 @@
 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.nodes import (
+    Call, Directive, IntrinsicCall, Loop, Routine, Schedule)
 from psyclone.psyir.transformations import (
-    ACCKernelsTrans, TransformationError, OMPTargetTrans)
+    ACCKernelsTrans, Matmul2CodeTrans, OMPTargetTrans, TransformationError)
 from psyclone.transformations import (
     LFRicColourTrans, LFRicOMPLoopTrans,
     LFRicRedundantComputationTrans, OMPParallelTrans,
@@ -59,9 +60,38 @@
 INVOKE_EXCLUSIONS = [
 ]
 
+# We won't attempt to inline calls to routines with names that contain
+# these strings (because they're not computationally important).
+INLINE_EXCLUSIONS = ["abort", "logging"]
+
 OFFLOAD_DIRECTIVES = os.getenv('LFRIC_OFFLOAD_DIRECTIVES', "none")
 
 
+def _replace_matmuls(sched: Schedule):
+    '''
+    Attempts to replace all MATMUL intrinsic calls with inline
+    code.
+
+    :param sched: schedule to transform.
+
+    '''
+    matrans = Matmul2CodeTrans()
+
+    for call in sched.walk(Call):
+        call: Call
+        # The NVIDIA compiler (as at 25.3) will sometimes fail to compile
+        # code with calls to MATMUL with a claim that they are not
+        # available on the device, e.g.:
+        # Call to NVHPC runtime function not supported -
+        #   pgf90_matmul_real4_i8
+        # Therefore, if we are unable to replace a MATMUL by generic code,
+        # the resulting TransformationError will signal (to the calling
+        # routine) that we are not to mark this kernel for offload.
+        if (isinstance(call, IntrinsicCall) and
+                call.intrinsic == IntrinsicCall.Intrinsic.MATMUL):
+            matrans.apply(call)
+
+
 def trans(psyir):
     '''Applies PSyclone colouring and GPU offloading transformations. Any
     kernels that cannot be offloaded to GPU are parallelised using OpenMP
@@ -106,7 +136,7 @@ def trans(psyir):
 
     for subroutine in psyir.walk(Routine):
 
-        print("Transforming invoke '{0}' ...".format(subroutine.name))
+        print(f"Transforming invoke '{subroutine.name}' ...")
 
         # Make setval_* compute redundantly to the level 1 halo if it
         # is in its own loop
@@ -122,7 +152,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
@@ -133,29 +164,37 @@ def trans(psyir):
                         const.VALID_DISCONTINUOUS_NAMES):
                     ctrans.apply(loop)
 
-        # Mark Kernels inside the loops over cells as GPU-enabled
-        # (alternatively we could inline them)
+        # Module-inline the Kernels inside the loops over cells and then mark
+        # them as GPU-enabled.
+        # (The latter step won't be necessary if/when we fully inline them.)
         for loop in subroutine.loops():
             if loop.iteration_space.endswith("cell_column"):
                 if offload:
                     for kern in loop.kernels():
+                        # Attempt to module-inline the kernel.
                         try:
                             mod_inline_trans.apply(kern)
                             print(f"Module-inlined kernel '{kern.name}'")
                         except TransformationError as err:
-                            print(f"Failed to module-inline '{kern.name}' due "
-                                  f"to:\n{err.value}")
+                            failed_inline.add(kern.name.lower())
+                            print(f"Failed to module-inline kernel "
+                                  f"'{kern.name}' due to:\n{err.value}")
                         try:
+                            # Ensure any MATMULs within the kernel are
+                            # replaced.
+                            for routine in kern.get_callees():
+                                _replace_matmuls(routine)
+                            # Finally, annotate the kernel routine for GPU.
                             gpu_annotation_trans.apply(kern)
                             print(f"Annotated kernel '{kern.name}'")
                         except TransformationError as err:
                             failed_to_offload.add(kern.name.lower())
                             print(f"Failed to annotate '{kern.name}' with "
                                   f"GPU-enabled directive due to:\n"
                                   f"{err.value}")
-                        # For annotated or inlined kernels we could attempt to
-                        # provide compile-time dimensions for the temporary
-                        # arrays and convert to code unsupported intrinsics.
+                        # For annotated/inlined kernels we could attempt to
+                        # provide compile-time dimensions for temporary arrays
+                        # and convert to code any unsupported intrinsics.
 
         # Add GPU offloading to loops unless they are over colours or are null.
         for loop in subroutine.walk(Loop):

examples/lfric/scripts/kernel_print.py

@@ -55,11 +55,11 @@ 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)
+            for ksched in kernel.get_callees():
+                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_callees()
+        # IGNORE polymorphic routines.
+        schedule = schedules[0]
         for child in schedule.children:
             exec_statements.append(self._visit(child))
         return exec_statements