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,14 +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, ProfileTrans,
     OMPMinimiseSyncTrans, Reference2ArrayRangeTrans,
-    ScalarisationTrans, IncreaseRankLoopArraysTrans, MaximalRegionTrans)
-from psyclone.transformations import TransformationError
+    ScalarisationTrans, IncreaseRankLoopArraysTrans, MaximalRegionTrans,
+    DataNodeToTempTrans, TransformationError)
 
 # USE statements to chase to gather additional symbol information.
 NEMO_MODULES_TO_IMPORT = [
@@ -526,3 +527,17 @@ 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":
+            arg = call.arguments[1]
+            dtype = arg.datatype
+            if isinstance(dtype, ArrayType) and isinstance(arg, Operation):
+                try:
+                    DataNodeToTempTrans().apply(arg)
+                except TransformationError:
+                    pass

src/psyclone/psyir/transformations/datanode_to_temp_trans.py

@@ -147,9 +147,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 +178,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 +250,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