Changes to OpenMP scripts to extract arguments from iom_put

changelog

@@ -1,3 +1,6 @@
+   13) PR #3373 towards #3412. Extract argument expressions from NEMO iom_put
+   calls in order to parallelise/offload their implicit loops.
+
    12) PR #3408 for #2812. Updates the Min/Max to code transformations so
    that they use the datatype of their arguments rather than assuming real.
    Also skips treesitter tests for Python < 3.10.

examples/nemo/scripts/omp_cpu_trans.py

@@ -117,7 +117,8 @@ def trans(psyir):
                 loopify_array_intrinsics=True,
                 convert_range_loops=True,
                 hoist_expressions=False,
-                scalarise_loops=False
+                scalarise_loops=False,
+                hoist_argument_expressions=True
         )
 
         if psyir.name not in PARALLELISATION_ISSUES:

examples/nemo/scripts/omp_gpu_trans.py

@@ -197,7 +197,8 @@ def trans(psyir):
                 loopify_array_intrinsics=True,
                 convert_range_loops=True,
                 increase_array_ranks=not NEMOV4,
-                hoist_expressions=True
+                hoist_expressions=True,
+                hoist_argument_expressions=True,
         )
         # Perform module-inlining of called routines.
         if INLINING_ENABLED:

examples/nemo/scripts/utils.py

@@ -41,15 +41,15 @@
 from psyclone.domain.common.transformations import KernelModuleInlineTrans
 from psyclone.psyir.nodes import (
     Assignment, Loop, Directive, Node, Reference, CodeBlock, Call,
-    Routine, Schedule, IntrinsicCall, StructureReference, IfBlock)
-from psyclone.psyir.symbols import DataSymbol
+    Routine, Schedule, IntrinsicCall, StructureReference, IfBlock,
+    Operation)
+from psyclone.psyir.symbols import DataSymbol, ArrayType
 from psyclone.psyir.transformations import (
     ArrayAssignment2LoopsTrans, HoistLoopBoundExprTrans, HoistLocalArraysTrans,
     HoistTrans, InlineTrans, Maxval2LoopTrans, Sum2LoopTrans, Minval2LoopTrans,
     Product2LoopTrans, ProfileTrans, OMPMinimiseSyncTrans,
     Reference2ArrayRangeTrans, ScalarisationTrans, IncreaseRankLoopArraysTrans,
-    MaximalRegionTrans)
-from psyclone.transformations import TransformationError
+    MaximalRegionTrans, TransformationError, DataNodeToTempTrans)
 
 # USE statements to chase to gather additional symbol information.
 NEMO_MODULES_TO_IMPORT = [
@@ -181,6 +181,7 @@ def normalise_loops(
         scalarise_loops: bool = False,
         increase_array_ranks: bool = False,
         hoist_expressions: bool = True,
+        hoist_argument_expressions: bool = True,
         ):
     ''' Normalise all loops in the given schedule so that they are in an
     appropriate form for the Parallelisation transformations to analyse
@@ -201,7 +202,14 @@ def normalise_loops(
         arrays.
     :param hoist_expressions: whether to hoist bounds and loop invariant
         statements out of the loop nest.
+    :param hoist_argument_expressions: whether to hoist array expressions
+        out of the containing Call.
     '''
+    # TODO #3412: This is currently limited to iom_put, we want to expand it
+    # throughout the code
+    if hoist_argument_expressions:
+        iom_put_argument_to_temporary(schedule.walk(Call))
+
     if hoist_local_arrays and schedule.name not in CONTAINS_STMT_FUNCTIONS:
         # Apply the HoistLocalArraysTrans when possible, it cannot be applied
         # to files with statement functions because it will attempt to put the
@@ -530,3 +538,25 @@ def _satisfies_minimum_region_rules(self, region: list[Node]) -> bool:
     routine_name = parent_routine.name if parent_routine else ""
     if routine_name not in PROFILING_IGNORE:
         MaximalProfilingOutsideDirectivesTrans().apply(children)
+
+
+def iom_put_argument_to_temporary(calls: list[Call]):
+    '''Extracts the second argument of all iom_put calls and puts them
+     in a temporary if they are an Operation with an array datatype.
+
+    :param calls: The list of calls in a subroutine whose arguments
+        may be moved into temporary storage to allow additional potential
+        parallelisation.
+
+     '''
+    for call in calls:
+        if call.symbol.name == "iom_put":
+            for arg in call.arguments:
+                dtype = arg.datatype
+                if (isinstance(dtype, ArrayType) and
+                    (isinstance(arg, Operation) or
+                        isinstance(arg, IntrinsicCall))):
+                    try:
+                        DataNodeToTempTrans().apply(arg, verbose=True)
+                    except TransformationError:
+                        pass

src/psyclone/psyir/frontend/fparser2.py

@@ -1112,8 +1112,8 @@ def _fparser2_tree_from_fparser2_reader(
             except (FortranSyntaxError, NoMatchError) as err:
                 raise ValueError(
                     f"Failed to parse the provided source code:\n{source_code}"
-                    "\nError was: {err}\nIs the input valid Fortran (note that"
-                    f" CPP directives must be handled by a pre-processor)?"
+                    f"\nError was: {err}\nIs the input valid Fortran (note "
+                    f"that CPP directives must be handled by a pre-processor)?"
                 ) from err
         try:
             # If it reaches this point a partial_code was provided, attempt

src/psyclone/psyir/nodes/intrinsic_call.py

@@ -439,6 +439,8 @@ def _iparity_return_type(node: IntrinsicCall) -> DataType:
 
     :returns: the computed datatype for the IntrinsicCall.
     """
+    # TODO #3415: Replace with _type_of_named_arg_accounting_for_dim_arg(
+    # node, "array").
     dtype = ScalarType(
         node.argument_by_name("array").datatype.intrinsic,
         node.argument_by_name("array").datatype.precision,
@@ -565,8 +567,13 @@ def _maxval_return_type(node: IntrinsicCall) -> DataType:
 
     :returns: the computed datatype for the IntrinsicCall.
     """
+    # TODO #3415: Replace with _type_of_named_arg_accounting_for_dim_arg(
+    # node, "array").
+    dtype = ScalarType(
+        node.argument_by_name("array").datatype.intrinsic,
+        node.argument_by_name("array").datatype.precision
+    )
     arg = node.argument_by_name("array")
-    dtype = arg.datatype.elemental_type
     if "dim" not in node.argument_names:
         return dtype
     # We have a dimension specified. We don't know the resultant shape
@@ -588,8 +595,8 @@ def _dot_product_return_type(node: IntrinsicCall) -> DataType:
     from psyclone.psyir.tools.type_info_computation import (
         compute_scalar_type
     )
-    veca_datatype = node.argument_by_name("vector_a").datatype
-    vecb_datatype = node.argument_by_name("vector_b").datatype
+    veca_datatype = node.argument_by_name("vector_a").datatype.elemental_type
+    vecb_datatype = node.argument_by_name("vector_b").datatype.elemental_type
     return compute_scalar_type(
         [ScalarType(
             veca_datatype.intrinsic, veca_datatype.precision
@@ -3215,7 +3222,8 @@ class Intrinsic(IAttr, Enum):
             optional_args={"kind": DataNode},
             return_type=lambda node: (
                 _type_of_scalar_with_optional_kind(
-                    node, node.argument_by_name("l").datatype.intrinsic,
+                    node,
+                    node.argument_by_name("l").datatype.intrinsic,
                     "kind",
                 ) if "kind" in node.argument_names else
                 _type_of_named_argument(node, "l")
@@ -4897,7 +4905,20 @@ def datatype(self) -> DataType:
         if isinstance(self.intrinsic.return_type, Callable):
             try:
                 return self.intrinsic.return_type(self)
+            except TypeError as err:
+                # If we get an invalid argument to a ScalarType constructor it
+                # means we attempted to pass either an UnresolvedType into the
+                # datatype
+                if ("ScalarType expected 'intrinsic' argument to be of type "
+                    in str(err)
+                    or "ScalarType expected 'precision' argument to be of "
+                        "type " in str(err)):
+                    return UnresolvedType()
+                # This should never happen, propogate as an InternalError.
+                outerr = err
             except AttributeError as err:
+                # This is to handle when we call .intrinsic or
+                # .precision on an UnresolvedType
                 # If we get an attribute error, and its because of attempting
                 # to lookup the precision or intrinsic, then it is likely
                 # due to looking up the datatype elements of an Unresolved
@@ -4908,13 +4929,15 @@ def datatype(self) -> DataType:
                     and "NoneType" not in
                         str(err)):
                     return UnresolvedType()
-                # Can't use debug string due to this being a potentially
-                # incomplete IntrinsicCall
-                raise InternalError(
-                    f"Failed to compute the datatype of a "
-                    f"'{self.intrinsic.name}' intrinsic. This is likely due "
-                    f"to not fully initialising the intrinsic correctly."
-                ) from err
+                outerr = err
+            # Fall through to the internalerror.
+            # Can't use debug string due to this being a potentially
+            # incomplete IntrinsicCall
+            raise InternalError(
+                f"Failed to compute the datatype of a "
+                f"'{self.intrinsic.name}' intrinsic. This is likely due "
+                f"to not fully initialising the intrinsic correctly."
+            ) from outerr
         else:
             return self.intrinsic.return_type