Changes to OpenMP scripts to extract arguments from iom_put

examples/nemo/scripts/omp_cpu_trans.py

@@ -40,8 +40,9 @@
 import os
 from utils import (
     insert_explicit_loop_parallelism, normalise_loops, add_profiling,
+    iom_put_argument_to_temporary,
     PARALLELISATION_ISSUES, NEMO_MODULES_TO_IMPORT)
-from psyclone.psyir.nodes import Routine
+from psyclone.psyir.nodes import Routine, Call
 from psyclone.transformations import OMPLoopTrans
 
 # Enable the insertion of profiling hooks during the transformation script
@@ -107,6 +108,10 @@ def trans(psyir):
     for subroutine in psyir.walk(Routine):
         print(f"Adding OpenMP threading to subroutine: {subroutine.name}")
 
+        # Extract any array operations from iom_put calls to temporary
+        # expressions that can be parallelised.
+        iom_put_argument_to_temporary(subroutine.walk(Call))
+
         if PROFILING_ENABLED:
             add_profiling(subroutine.children)
 

examples/nemo/scripts/omp_gpu_trans.py

@@ -40,8 +40,9 @@
 import os
 from utils import (
     add_profiling, inline_calls, insert_explicit_loop_parallelism,
-    normalise_loops, PARALLELISATION_ISSUES, NEMO_MODULES_TO_IMPORT)
-from psyclone.psyir.nodes import Routine, Loop
+    normalise_loops, iom_put_argument_to_temporary,
+    PARALLELISATION_ISSUES, NEMO_MODULES_TO_IMPORT)
+from psyclone.psyir.nodes import Routine, Loop, Call
 from psyclone.psyir.transformations import (
     OMPTargetTrans, OMPDeclareTargetTrans)
 from psyclone.transformations import (
@@ -199,6 +200,10 @@ def trans(psyir):
             if "pp_len" not in symtab:
                 symtab.add(symtab.lookup("pp_len"))
 
+        # Extract any array operations from iom_put calls to temporary
+        # expressions that can be parallelised.
+        iom_put_argument_to_temporary(subroutine.walk(Call))
+
         normalise_loops(
                 subroutine,
                 hoist_local_arrays=False,

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 = [
@@ -530,3 +530,22 @@ 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.'''
+    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)):
+                    try:
+                        DataNodeToTempTrans().apply(arg)
+                    except TransformationError as err:
+                        call.append_preceding_comment(
+                            f"Couldn't pull the argument {arg} to a "
+                            f"temporary due to the following error: "
+                            f"{str(err.value)}"
+                        )

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

@@ -565,8 +565,11 @@ def _maxval_return_type(node: IntrinsicCall) -> DataType:
 
     :returns: the computed datatype for the IntrinsicCall.
     """
+    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 +591,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 +3218,9 @@ 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")
@@ -3763,8 +3768,10 @@ class Intrinsic(IAttr, Enum):
             optional_args={"vector": DataNode},
             return_type=lambda node: ArrayType(
                 ScalarType(
-                    node.argument_by_name("array").datatype.intrinsic,
-                    node.argument_by_name("array").datatype.precision),
+                    node.argument_by_name("array").datatype.
+                    intrinsic,
+                    node.argument_by_name("array").datatype.
+                    precision),
                 [ArrayType.Extent.DEFERRED]
             ),
             reference_accesses=lambda node: (
@@ -4444,8 +4451,10 @@ class Intrinsic(IAttr, Enum):
             optional_args={},
             return_type=lambda node: ArrayType(
                 ScalarType(
-                    node.argument_by_name("source").datatype.intrinsic,
-                    node.argument_by_name("source").datatype.precision),
+                    node.argument_by_name("source").datatype.
+                    intrinsic,
+                    node.argument_by_name("source").datatype.
+                    precision),
                 ([ArrayType.Extent.DEFERRED] *
                  (len(node.argument_by_name("source").datatype.shape) + 1)
                  if isinstance(node.argument_by_name("source").datatype,
@@ -4714,8 +4723,10 @@ class Intrinsic(IAttr, Enum):
                                 ArrayType))
                 else ArrayType(
                     ScalarType(
-                        node.argument_by_name("mold").datatype.intrinsic,
-                        node.argument_by_name("mold").datatype.precision
+                        node.argument_by_name("mold").datatype.
+                        intrinsic,
+                        node.argument_by_name("mold").datatype.
+                        precision
                     ),
                     [ArrayType.Extent.DEFERRED])
             ),
@@ -4737,8 +4748,10 @@ class Intrinsic(IAttr, Enum):
                 arg_names=(("matrix",),)),
             optional_args={},
             return_type=lambda node: ArrayType(ScalarType(
-                node.argument_by_name("matrix").datatype.intrinsic,
-                node.argument_by_name("matrix").datatype.precision),
+                node.argument_by_name("matrix").datatype.
+                intrinsic,
+                node.argument_by_name("matrix").datatype.
+                precision),
                 [node.argument_by_name("matrix").datatype.shape[1],
                  node.argument_by_name("matrix").datatype.shape[0]]
             ),
@@ -4897,7 +4910,21 @@ 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()
+                # Is this reachable? Tested via monkeypatch as there may be
+                # some edge case I can't think of.
+                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 +4935,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
 

src/psyclone/psyir/transformations/datanode_to_temp_trans.py

@@ -42,10 +42,14 @@
     Assignment,
     Call,
     DataNode,
+    IfBlock,
     IntrinsicCall,
+    Loop,
     Range,
     Reference,
     Statement,
+    Schedule,
+    UnaryOperation,
 )
 from psyclone.psyir.symbols.datatypes import (
     ArrayType,
@@ -147,9 +151,11 @@ def validate(self, node: DataNode, **kwargs):
                     symbols.update(element.upper.get_all_accessed_symbols())
                 # Compare the symbols in the array bounds with the symbols
                 # already in the scope.
-                scope_symbols = node.scope.symbol_table.get_symbols()
+                scope_table = node.scope.symbol_table
                 for sym in symbols:
-                    scoped_name_sym = scope_symbols.get(sym.name, None)
+                    scoped_name_sym = scope_table.lookup(
+                        sym.name, otherwise=None
+                    )
                     # If sym is not scoped_name_sym, then there is a
                     # symbol collision from an imported symbol.
                     if scoped_name_sym and sym is not scoped_name_sym:
@@ -176,9 +182,9 @@ def validate(self, node: DataNode, **kwargs):
                     # If its an imported symbol we need to check if its
                     # the same import interface.
                     if isinstance(sym.interface, ImportInterface):
-                        scoped_name_sym = scope_symbols.get(
-                                sym.interface.container_symbol.name,
-                                None
+                        scoped_name_sym = scope_table.lookup(
+                            sym.interface.container_symbol.name,
+                            otherwise=None
                         )
                         if scoped_name_sym and not isinstance(
                                 scoped_name_sym, ContainerSymbol):
@@ -248,18 +254,20 @@ def apply(self, node: DataNode, storage_name: str = "", **kwargs):
                     symbols.update(element.lower.get_all_accessed_symbols())
                 if isinstance(element.upper, DataNode):
                     symbols.update(element.upper.get_all_accessed_symbols())
-                scope_symbols = node.scope.symbol_table.get_symbols()
+                scope_table = node.scope.symbol_table
                 for sym in symbols:
-                    scoped_name_sym = scope_symbols.get(sym.name, None)
+                    scoped_name_sym = scope_table.lookup(
+                        sym.name, otherwise=None
+                    )
                     # If no symbol with the name exists then create one.
                     if not scoped_name_sym:
                         sym_copy = sym.copy()
                         if isinstance(sym_copy.interface, ImportInterface):
                             # Check if the ContainerSymbol is already in the
                             # interface
-                            container = scope_symbols.get(
+                            container = scope_table.lookup(
                                 sym_copy.interface.container_symbol.name,
-                                None
+                                otherwise=None
                             )
                             if container is None:
                                 # Add the container symbol to the symbol table
@@ -300,10 +308,11 @@ def apply(self, node: DataNode, storage_name: str = "", **kwargs):
         # Create a Reference to the new symbol
         new_ref = Reference(symbol)
 
-        # Find the parent and position of the statement containing the
-        # DataNode.
-        parent = node.ancestor(Statement).parent
-        pos = node.ancestor(Statement).position
+        # Find the containing schedule and position of the statement
+        # containing the DataNode.
+        schedule = node.ancestor(Schedule)
+        path = node.path_from(schedule)
+        pos = path[0]
 
         # Replace the datanode with the new reference
         node.replace_with(new_ref)
@@ -312,10 +321,10 @@ def apply(self, node: DataNode, storage_name: str = "", **kwargs):
         assign = Assignment.create(new_ref.copy(), node)
 
         # Add the assignment into the tree.
-        parent.addchild(assign, pos)
+        schedule.addchild(assign, pos)
 
         # If the datatype is an array, we need to allocate the array
-        # before the statement too.
+        # before the statement too if its not already allocated.
         if isinstance(datatype, ArrayType):
             # Create an array reference to the symbol with the dimensions
             # returned by the datatype call earlier.
@@ -329,9 +338,40 @@ def apply(self, node: DataNode, storage_name: str = "", **kwargs):
                 IntrinsicCall.Intrinsic.ALLOCATE,
                 (ref,)
             )
-            # Add the allocate statement into the tree immediately before
-            # its use.
-            parent.addchild(intrinsic, pos)
+            allocated = IntrinsicCall.create(
+                IntrinsicCall.Intrinsic.ALLOCATED,
+                (Reference(symbol),)
+            )
+
+            ifblock = IfBlock.create(
+                UnaryOperation.create(
+                    UnaryOperation.Operator.NOT,
+                    allocated),
+                [intrinsic]
+            )
+            # If the shape doesn't contain array references then we can hoist
+            # the allocate statement outside of any ancestor loops.
+            hoistable = True
+            for shape in ref.indices:
+                for ref2 in shape.walk(Reference):
+                    if isinstance(ref2, ArrayReference):
+                        hoistable = False
+            # If we can hoist the allocate, find the highest level Loop
+            # ancestor and set the schedule and position to place the
+            # allocate before this loop.
+            if hoistable:
+                loop_anc = schedule.ancestor(Loop)
+                finger = loop_anc
+                while finger:
+                    loop_anc = finger
+                    finger = finger.ancestor(Loop)
+                if loop_anc:
+                    pos = loop_anc.position
+                    schedule = loop_anc.ancestor(Schedule)
+
+            # Add the allocate statement and the containing ifblock into the
+            # tree immediately before its use.
+            schedule.addchild(ifblock, pos)
 
 
 __all__ = ["DataNodeToTempTrans"]