Generalise hoist_arguments_to_temporaries in NEMO utils to all arguments with array operations

changelog

@@ -1,3 +1,6 @@
+   24) PR #3429 for #3412. Generalise hoist_arguments_to_temporaries in NEMO utils
+   to all arguments with array operations.
+
    23) PR #3436 for #3135. Fixes some cyclic dependencies by replacing
    the ScalarType import-time instantiations.
 

examples/nemo/scripts/utils.py

@@ -43,7 +43,7 @@
     Assignment, Loop, Directive, Node, Reference, CodeBlock, Call,
     Routine, Schedule, IntrinsicCall, StructureReference, IfBlock,
     Operation)
-from psyclone.psyir.symbols import DataSymbol, ArrayType
+from psyclone.psyir.symbols import DataSymbol, ArrayType, ScalarType
 from psyclone.psyir.transformations import (
     ArrayAssignment2LoopsTrans, HoistLoopBoundExprTrans, HoistLocalArraysTrans,
     HoistTrans, InlineTrans, Maxval2LoopTrans, Sum2LoopTrans, Minval2LoopTrans,
@@ -205,10 +205,6 @@ def normalise_loops(
     :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
@@ -280,6 +276,20 @@ def normalise_loops(
                 except TransformationError:
                     pass
 
+    if hoist_argument_expressions:
+        hoist_arguments_to_temporaries(schedule.walk(Call))
+        normalise_loops(
+            schedule,
+            hoist_local_arrays=hoist_local_arrays,
+            convert_array_notation=convert_array_notation,
+            loopify_array_intrinsics=loopify_array_intrinsics,
+            convert_range_loops=convert_range_loops,
+            scalarise_loops=scalarise_loops,
+            increase_array_ranks=False,
+            hoist_expressions=hoist_expressions,
+            # Make sure we never repeat this step.
+            hoist_argument_expressions=False,
+        )
     # TODO #1928: In order to perform better on the GPU, nested loops with two
     # sibling inner loops need to be fused or apply loop fission to the
     # top level. This would allow the collapse clause to be applied.
@@ -541,23 +551,30 @@ def _satisfies_minimum_region_rules(self, region: list[Node]) -> bool:
         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.
+def hoist_arguments_to_temporaries(calls: list[Call]):
+    '''Extracts the arguments of all calls and puts them
+     in a temporary result 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
+        for arg in call.arguments:
+            dtype = arg.datatype
+            # Only extract expressions that can potentially be loopfied,
+            # i.e. operations over arrays.
+            if (isinstance(dtype, ArrayType) and
+                (isinstance(arg, Operation) or
+                    isinstance(arg, IntrinsicCall))):
+                try:
+                    # TODO #3443: Does the fixes in that PR fix this, or
+                    # maybe the DataNodeToTempTrans has to skip characters
+                    if (isinstance(dtype.elemental_type, ScalarType)
+                        and dtype.elemental_type.intrinsic ==
+                            ScalarType.Intrinsic.CHARACTER):
+                        continue
+                    DataNodeToTempTrans().apply(arg, verbose=True)
+                except TransformationError:
+                    pass

src/psyclone/psyir/transformations/datanode_to_temp_trans.py

@@ -48,6 +48,7 @@
     Loop,
     Range,
     Reference,
+    Routine,
     Statement,
     Schedule,
     UnaryOperation,
@@ -230,7 +231,10 @@ def validate(self, node: DataNode, **kwargs):
                 f"Statement node which is not supported."
             )
 
-        if isinstance(dtype, (UnresolvedType, UnsupportedFortranType)):
+        if (isinstance(dtype, (UnresolvedType, UnsupportedFortranType))
+                or (isinstance(dtype, ArrayType) and
+                    isinstance(dtype.elemental_type,
+                               (UnresolvedType, UnsupportedFortranType)))):
             failing_symbols = []
             symbols = node.get_all_accessed_symbols()
             for sym in symbols:
@@ -343,14 +347,19 @@ def apply(self, node: DataNode, storage_name: str = "",
                                       allocatable_datatype.shape])
 
         # Create a symbol of the relevant type.
+        containing_routine = node.ancestor(Routine)
+        if containing_routine:
+            sym_tab = containing_routine.symbol_table
+        else:
+            sym_tab = node.scope.symbol_table
         if not storage_name:
-            symbol = node.scope.symbol_table.new_symbol(
+            symbol = sym_tab.new_symbol(
                 root_name="tmp",
                 symbol_type=DataSymbol,
                 datatype=datatype
             )
         else:
-            symbol = node.scope.symbol_table.new_symbol(
+            symbol = sym_tab.new_symbol(
                 root_name=storage_name,
                 symbol_type=DataSymbol,
                 datatype=datatype

src/psyclone/psyir/transformations/intrinsics/array_reduction_base_trans.py

@@ -126,7 +126,6 @@ def validate(self, node, options=None, **kwargs):
             raise TransformationError(
                 f"The dimension argument to {self._INTRINSIC_NAME} is not "
                 f"yet supported.")
-
         if isinstance(node.datatype, (UnresolvedType, UnsupportedType)):
             raise TransformationError(
                 f"Error in {self.name} transformation. Cannot create "
@@ -178,6 +177,9 @@ def apply(self, node, options=None, verbose: bool = False, **kwargs):
         :param options: options for the transformation.
         :type options: Optional[Dict[str, Any]]
 
+        :raises TransformationError: if node only contains Reference nodes,
+            and they can't be convered to ArrayReferences by the calls to
+            Reference2ArrayRangeTrans.
         '''
         # TODO 2668: options are now deprecated:
         if options:
@@ -221,6 +223,11 @@ def apply(self, node, options=None, verbose: bool = False, **kwargs):
         # resulting in the following code being created:
         # a(:,:) = a(:,:)+b(:,:)
         array_refs = new_array_expr.walk(ArrayReference)
+        if len(array_refs) == 0:
+            raise TransformationError(
+                f"Can't apply {self.name} to {node.debug_string()} due "
+                f"to no ArrayReference nodes present."
+            )
         assignment = Assignment.create(array_refs[0].copy(),
                                        new_array_expr.detach())
 

src/psyclone/psyir/transformations/reference2arrayrange_trans.py

@@ -232,7 +232,6 @@ def apply(self, node, options=None, **kwargs):
 
         '''
         self.validate(node, **kwargs)
-
         # The following cases do not need expansions
         if node.parent and isinstance(node.parent, Call):
             if node is node.parent.routine:

src/psyclone/tests/psyir/transformations/datanode_to_temp_trans_test.py

@@ -40,7 +40,7 @@
 from psyclone.configuration import Config
 from psyclone.psyir.frontend.fortran import FortranReader
 from psyclone.psyir.nodes import (
-    Assignment, Reference
+    Assignment, Loop, Reference, Routine
 )
 from psyclone.psyir.symbols import (
     DataSymbol, ScalarType
@@ -299,6 +299,9 @@ def test_datanodetotemptrans_apply(fortran_reader, fortran_writer, tmp_path):
     psyir = fortran_reader.psyir_from_source(code)
     assign = psyir.walk(Assignment)[0]
     dtrans.apply(assign.rhs.operands[1])
+    # Check the tmp symbol is at the routine scope.
+    routine = psyir.walk(Routine)[0]
+    assert "tmp" in routine.symbol_table
     out = fortran_writer(psyir)
     assert """  integer, allocatable, dimension(:,:) :: tmp
 
@@ -309,6 +312,27 @@ def test_datanodetotemptrans_apply(fortran_reader, fortran_writer, tmp_path):
   d = c + tmp""" in out
     assert Compile(tmp_path).string_compiles(out)
 
+    # Check the symbol is still created if the assignment is in some other
+    # non-routine scope.
+    code = """subroutine test()
+    integer :: a, b, c
+    integer :: i
+
+    do i = 1, 100
+      a = b + c
+    end do
+    end subroutine test"""
+    psyir = fortran_reader.psyir_from_source(code)
+    loop = psyir.walk(Loop)[0].detach()
+    assign = loop.walk(Assignment)[0]
+    dtrans.apply(assign.rhs)
+    assert "tmp" in assign.scope.symbol_table
+    out = fortran_writer(loop)
+    assert """do i = 1, 100, 1
+  tmp = b + c
+  a = tmp
+enddo""" in out
+
     code = """subroutine test()
         real :: a
         integer :: b

src/psyclone/tests/psyir/transformations/intrinsics/sum2loop_trans_test.py

@@ -141,3 +141,59 @@ def test_apply(fortran_reader, fortran_writer, tmpdir):
     result = fortran_writer(psyir)
     assert expected in result
     assert Compile(tmpdir).string_compiles(result)
+
+
+def test_nested_sums_error_cases(fortran_reader):
+    '''Test that the transformation fails correctly if we have nested sums
+    as the reference can't be converted to an arrayreference as the reference
+    is inside a non-elemental function.'''
+    code = """subroutine sum_test()
+    integer :: n, m
+    real , dimension(:, :) :: array
+    real :: result
+
+    result = sum(sum(array, dim=2)) * array(n,m)
+    end subroutine"""
+    psyir = fortran_reader.psyir_from_source(code)
+    intrinsic_node = psyir.children[0].children[0].rhs.children[0]
+    trans = Sum2LoopTrans()
+    with pytest.raises(TransformationError) as err:
+        trans.apply(intrinsic_node)
+    assert ("Can't apply Sum2LoopTrans to SUM(SUM(array, 2)) due "
+            "to no ArrayReference nodes present." in str(err.value))
+
+    code = """subroutine sum_test()
+    integer :: n, m
+    real, dimension(:, :) :: array
+    real :: result
+
+    result = sum(sum(array  + array(:,:), dim=2))
+    end subroutine"""
+    psyir = fortran_reader.psyir_from_source(code)
+    intrinsic_node = psyir.children[0].children[0].rhs
+    trans = Sum2LoopTrans()
+    with pytest.raises(TransformationError) as err:
+        trans.apply(intrinsic_node)
+    assert ("Transformation Error: ArrayAssignment2LoopsTrans does not "
+            "support statements containing dependencies that would generate "
+            "loop-carried dependencies when naively converting them to a "
+            "loop, but found" in str(err.value))
+
+    code = """subroutine sum_test()
+    integer :: n, m
+    real, dimension(:, :) :: array
+    integer, dimension(1) :: dimensions
+    real :: result
+    result = sum(sum(array, dim=dimensions(2)))
+    end subroutine"""
+    psyir = fortran_reader.psyir_from_source(code)
+    intrinsic_node = psyir.children[0].children[0].rhs
+    trans = Sum2LoopTrans()
+    with pytest.raises(TransformationError) as err:
+        trans.apply(intrinsic_node)
+    assert ("Transformation Error: Error in ArrayAssignment2LoopsTrans "
+            "transformation. The LHS of the supplied Assignment node "
+            "should contain an array accessor with at least one of its "
+            "dimensions being a Range, but none were found in "
+            "'dimensions(2) = SUM(array, dimensions(2))"
+            in str(err.value))