(closes #2884) Allow WHERE with imported symbols by changing Ref2ArrayRange behaviour to do the validations

examples/nemo/scripts/utils.py

@@ -212,17 +212,11 @@ def normalise_loops(
             pass
 
     if convert_array_notation:
-        # Make sure all array dimensions are explicit
         for reference in schedule.walk(Reference):
-            part_of_the_call = reference.ancestor(Call)
-            if part_of_the_call:
-                if not part_of_the_call.is_elemental:
-                    continue
-            if isinstance(reference.symbol, DataSymbol):
-                try:
-                    Reference2ArrayRangeTrans().apply(reference)
-                except TransformationError:
-                    pass
+            try:
+                Reference2ArrayRangeTrans().apply(reference)
+            except TransformationError:
+                pass
 
     if loopify_array_intrinsics:
         for intr in schedule.walk(IntrinsicCall):

src/psyclone/psyad/transformations/preprocess.py

@@ -78,6 +78,11 @@ def preprocess_trans(kernel_psyir, active_variable_names):
         try:
             reference2arrayrange_trans.apply(reference)
         except TransformationError:
+            # TODO #3269: For now we consider that everywhere this fails does
+            # not need array notation (e.g. it is a scalar or a non-elemental
+            # routine). This is unsafe, 'psyad' should not proceed if this
+            # happens, to make it sagfe we need to follow the necessary imports
+            # to guarantee the datatype of imported symbols.
             pass
 
     for call in kernel_psyir.walk(IntrinsicCall):

src/psyclone/psyir/frontend/fparser2.py

@@ -4238,31 +4238,22 @@ def _array_syntax_to_indexed(self, parent, loop_vars):
         # pylint: disable=import-outside-toplevel
         from psyclone.psyir.transformations import (
                 Reference2ArrayRangeTrans, TransformationError)
-        # Convert References to arrays to use the array range notation unless
-        # they have an IntrinsicCall parent.
+        # For each reference, we need to know if they are used as an array
+        # access (and convert it to ArrayReference) or not. If we cannot
+        # guarantee how it is used we cannot safely proceed.
         for ref in parent.walk(Reference):
-            if isinstance(ref.symbol.interface, (ImportInterface,
-                                                 UnresolvedInterface)):
+            try:
+                Reference2ArrayRangeTrans().apply(ref)
+            except TransformationError as error:
                 raise NotImplementedError(
-                        "PSyclone doesn't yet support reference to imported "
-                        "symbols inside WHERE clauses.")
-            call_ancestor = ref.ancestor(Call)
-            if (isinstance(ref.symbol, DataSymbol) and not call_ancestor):
-                try:
-                    Reference2ArrayRangeTrans().apply(ref)
-                except TransformationError:
-                    pass
-            elif (isinstance(ref.symbol, DataSymbol) and call_ancestor
-                  is not None and call_ancestor.is_elemental):
-                try:
-                    Reference2ArrayRangeTrans().apply(ref)
-                except TransformationError:
-                    pass
+                    f"WHERE not supported because '{ref.name}' cannot "
+                    f"be converted to an array due to: {error}")
         table = parent.scope.symbol_table
         one = Literal("1", INTEGER_TYPE)
         arrays = parent.walk(ArrayMixin)
 
         first_rank = None
+        first_array = None
         for array in arrays:
             # Check that this is a supported array reference and that
             # all arrays are of the same rank
@@ -4273,30 +4264,23 @@ def _array_syntax_to_indexed(self, parent, loop_vars):
                 # ignore it.
                 continue
 
-            # If it has a Call ancestor we need to check if its a
+            # If it is an argument to a Call we need to check if its a
             # non-elemental function, in which case we should skip
             # changing it.
-            call_ancestor = array.ancestor(Call)
-            if call_ancestor:
-                if call_ancestor.is_elemental is None:
-                    raise NotImplementedError(
-                        f"Found a function call inside a where clause with "
-                        f"unknown elemental status: "
-                        f"{call_ancestor.debug_string()}")
-                # If it is none-elemental, we leave this array reference as it
-                # is
-                if not call_ancestor.is_elemental:
+            if isinstance(array.parent, Call):
+                if not array.parent.is_elemental:
                     continue
-                # Otherwise, we continue replacing the range with the loop idx
 
             if first_rank:
                 if rank != first_rank:
                     raise NotImplementedError(
                         f"Found array sections of differing ranks within a "
-                        f"WHERE construct: array section of {array.name} has "
-                        f"rank {rank}")
+                        f"WHERE construct: array section of '{array.name}' "
+                        f" has rank '{rank}', but '{first_array.name}' has "
+                        f"rank '{first_rank}'")
             else:
                 first_rank = rank
+                first_array = array
 
             base_ref = _copy_full_base_reference(array)
             array_ref = array.ancestor(Reference, include_self=True)
@@ -4410,11 +4394,9 @@ def _contains_intrinsic_reduction(pnodes):
                 if (intr.children[0].string in
                         Fortran2003.Intrinsic_Name.array_reduction_names):
                     # These intrinsics are only a problem if they return an
-                    # array rather than a scalar.
-                    arg_specs = walk(intr.children[1],
-                                     Fortran2003.Actual_Arg_Spec)
-                    if any(spec.children[0].string == 'dim'
-                           for spec in arg_specs):
+                    # array rather than a scalar (which can only happen if
+                    # there is more than one argument).
+                    if len(intr.children[1].children) > 1:
                         return True
             return False
 
@@ -4481,27 +4463,14 @@ def _contains_intrinsic_reduction(pnodes):
         # We want to convert all the plain references that are arrays to use
         # explicit array syntax.
         # TODO 2722: Should have the same logic as array_syntax_to_indexed
-        # regarding UnresolvedInterface and Elemental calls?
         references = fake_parent.walk(Reference)
         for ref in references:
-            call_ancestor = ref.ancestor(Call)
-            elemental_ancestor = (call_ancestor is None or
-                                  call_ancestor.is_elemental)
-            # TODO 2884: We should be able to handle this imported symbol
-            # better. If we can, we need to handle a case where is_elemental
-            # can be None.
-            if isinstance(ref.symbol.interface, (ImportInterface,
-                                                 UnresolvedInterface)):
+            try:
+                Reference2ArrayRangeTrans().apply(ref)
+            except TransformationError as error:
                 raise NotImplementedError(
-                        f"PSyclone doesn't yet support references to imported/"
-                        f"unresolved symbols inside WHERE clauses: "
-                        f"'{ref.symbol.name}' is unresolved.")
-            if (isinstance(ref.symbol, DataSymbol) and
-                    elemental_ancestor):
-                try:
-                    Reference2ArrayRangeTrans().apply(ref)
-                except TransformationError:
-                    pass
+                    f"WHERE not supported because '{ref.name}' cannot "
+                    f"be converted to an array due to: {error}")
 
         arrays = fake_parent.walk(ArrayMixin)
 

src/psyclone/psyir/transformations/arrayassignment2loops_trans.py

@@ -118,10 +118,7 @@ def apply(self, node, options=None):
         # If there is any array reference without the accessor syntax,
         # we need to add it first.
         for reference in node.walk(Reference):
-            try:
-                Reference2ArrayRangeTrans().apply(reference)
-            except TransformationError:
-                pass
+            Reference2ArrayRangeTrans().apply(reference)
 
         # Start by the rightmost array range
         for lhs_range in reversed(node.lhs.walk(Range)):

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

@@ -242,10 +242,7 @@ def apply(self, node, options=None, **kwargs):
         # transformation.
         rhs_parent = rhs.parent
         for reference in rhs.walk(Reference):
-            try:
-                reference2arrayrange.apply(reference)
-            except TransformationError:
-                pass
+            reference2arrayrange.apply(reference)
         # Reset rhs from its parent as the previous transformation
         # makes the value of rhs become invalid. We know there is only
         # one child so can safely use children[0].
@@ -254,10 +251,7 @@ def apply(self, node, options=None, **kwargs):
             mask_ref_parent = mask_ref.parent
             mask_ref_index = mask_ref.position
             for reference in mask_ref.walk(Reference):
-                try:
-                    reference2arrayrange.apply(reference)
-                except TransformationError:
-                    pass
+                reference2arrayrange.apply(reference)
             mask_ref = mask_ref_parent.children[mask_ref_index]
 
         # Step 2: Put the intrinsic's extracted expression (stored in

src/psyclone/psyir/transformations/reference2arrayrange_trans.py

@@ -43,18 +43,31 @@
 '''
 from psyclone.errors import LazyString
 from psyclone.psyGen import Transformation
-from psyclone.psyir.nodes import (ArrayReference, Assignment, Call,
-                                  Literal, Range, Reference)
-from psyclone.psyir.symbols import INTEGER_TYPE
+from psyclone.psyir.nodes import (
+    ArrayReference, Call, Literal, Range, Reference,
+    StructureReference, ArrayOfStructuresReference, Assignment)
+from psyclone.psyir.symbols import (
+    INTEGER_TYPE, DataSymbol, UnresolvedType, UnsupportedType, ArrayType,
+    ScalarType)
 from psyclone.psyir.transformations.transformation_error import (
     TransformationError)
 from psyclone.utils import transformation_documentation_wrapper
 
 
 @transformation_documentation_wrapper
 class Reference2ArrayRangeTrans(Transformation):
-    '''Provides a transformation from PSyIR Array Notation (a reference to
-    an Array) to a PSyIR Range. For example:
+    '''
+    Transformation to convert plain References of array symbols to
+    ArrayReferances with full-extent ranges if it is semantically equivalent
+    to do so (e.g. it won't convert call arguments because it would change the
+    bounds values).
+
+    Note that if the provided node does not need to be modified (
+    e.g. a Reference to a scalar or an ArrayReference to an array), the
+    transformation will succeed. However, if we cannot guarantee the type of
+    the symbol, or the validity of the transformations (e.g. it is in a call
+    that we don't know if it is elemental or not), the transformation will
+    fail.
 
     >>> from psyclone.psyir.backend.fortran import FortranWriter
     >>> from psyclone.psyir.frontend.fortran import FortranReader
@@ -84,6 +97,7 @@ class Reference2ArrayRangeTrans(Transformation):
     structures, see issue #1858.
 
     '''
+
     def validate(self, node, options=None, **kwargs):
         '''Check that the node is a Reference node and that we have all
         information necessary to decide if it can be expanded.
@@ -103,28 +117,56 @@ def validate(self, node, options=None, **kwargs):
         super().validate(node, **kwargs)
         self.validate_options(**kwargs)
 
-        # TODO issue #1858. Add support for structures containing arrays.
+        if node and isinstance(node.parent, Call):
+            if node is node.parent.routine:
+                return
+            if node.parent.is_elemental is None:
+                raise TransformationError(LazyString(
+                    lambda: f"The supplied node is passed as an argument to a "
+                    f"Call that may or may not be elemental: "
+                    f"'{node.parent.debug_string().strip()}'. Consider "
+                    f"adding the function's filename to RESOLVE_IMPORTS."))
+            if not node.parent.is_elemental:
+                return
+        assignment = node.ancestor(Assignment) if node else None
+        if assignment and assignment.is_pointer:
+            raise TransformationError(
+                f"{type(self).__name__} can not be applied to references"
+                f" inside pointer assignments, but found '{node.name}' in"
+                f" {assignment.debug_string()}")
+
         # pylint: disable=unidiomatic-typecheck
+        if (
+            type(node) is ArrayReference or
+            type(node) is ArrayOfStructuresReference
+        ):
+            # If it is already an Array access, it does not need expansion
+            # nor further validation
+            return
+
+        if type(node) is StructureReference:
+            # TODO #1858: Add support for expansion of structures
+            return
+
         if not type(node) is Reference:
             raise TransformationError(
                 f"The supplied node should be a Reference but found "
                 f"'{type(node).__name__}'.")
-        if not node.symbol.is_array:
+        if (
+            not isinstance(node.symbol, DataSymbol) or
+            isinstance(node.symbol.datatype, (UnresolvedType,
+                                              UnsupportedType))
+        ):
+            if (
+                isinstance(node.symbol, DataSymbol) and
+                isinstance(node.symbol.datatype, UnsupportedType) and
+                isinstance(node.symbol.datatype.partial_datatype,
+                           (ScalarType, ArrayType))
+            ):
+                return
             raise TransformationError(
                 f"The supplied node should be a Reference to a symbol "
-                f"that is an array, but '{node.symbol.name}' is not.")
-        if isinstance(node.parent, Call) and not node.parent.is_elemental:
-            raise TransformationError(LazyString(
-                lambda: f"The supplied node is passed as an argument to a "
-                f"Call to a non-elemental routine ("
-                f"{node.parent.debug_string().strip()}) and should not be "
-                f"transformed."))
-        assignment = node.ancestor(Assignment)
-        if assignment and assignment.is_pointer:
-            raise TransformationError(
-                f"'{type(self).__name__}' can not be applied to references"
-                f" inside pointer assignments, but found '{node.name}' in"
-                f" {assignment.debug_string()}")
+                f"of known type, but '{node.symbol}' is not.")
 
     def apply(self, node, options=None, **kwargs):
         '''Apply the Reference2ArrayRangeTrans transformation to the specified
@@ -138,7 +180,22 @@ def apply(self, node, options=None, **kwargs):
         '''
         self.validate(node, **kwargs)
 
+        # We have validated that it is a reference to a symbol with a datatype
+        # that is not UnresolvedType or UnknownType
         symbol = node.symbol
+
+        # The following cases do not need expansions
+        # pylint: disable=unidiomatic-typecheck
+        if type(node) in [ArrayReference, ArrayOfStructuresReference]:
+            return
+        if node.parent and isinstance(node.parent, Call):
+            if node is node.parent.routine:
+                return
+            if not node.parent.is_elemental:
+                return
+        if not symbol.is_array:
+            return
+
         indices = []
         for idx, _ in enumerate(symbol.shape):
             lbound, ubound = symbol.get_bounds(idx)

src/psyclone/tests/psyad/transformations/test_preprocess.py

@@ -76,17 +76,19 @@ def test_preprocess_reference2arrayrange(tmpdir, fortran_reader,
     '''
     code = (
         "program test\n"
+        "use other, only: g\n"
         "real, dimension(10,10) :: a,b,c,e,f\n"
         "real, dimension(10) :: d\n"
         "integer :: i\n"
         "a = b * c\n"
         "do i = lbound(d,1), ubound(d,1)\n"
         "  d(i) = 0.0\n"
         "end do\n"
-        "e = f\n"
+        "e = f + g\n"
         "end program test\n")
     expected = (
         "program test\n"
+        "  use other, only : g\n"
         "  real, dimension(10,10) :: a\n"
         "  real, dimension(10,10) :: b\n"
         "  real, dimension(10,10) :: c\n"
@@ -104,13 +106,18 @@ def test_preprocess_reference2arrayrange(tmpdir, fortran_reader,
         "  do i = LBOUND(d, dim=1), UBOUND(d, dim=1), 1\n"
         "    d(i) = 0.0\n"
         "  enddo\n"
-        "  e(:,:) = f(:,:)\n\n"
+        "  e(:,:) = g + f(:,:)\n\n"
         "end program test\n")
     psyir = fortran_reader.psyir_from_source(code)
     preprocess_trans(psyir, ["a", "c"])
     result = fortran_writer(psyir)
     assert result == expected
-    assert Compile(tmpdir).string_compiles(result)
+    # TODO #3269: Currently this tests shows that psyad converts references to
+    # array_references even when there are imported symbols with unknown type
+    # ('g' in the example above). This demonstrate this we added an import that
+    # makes the test not compilable, this will be fixed when psyad can follow
+    # dependencies.
+    # assert Compile(tmpdir).string_compiles(result)
 
 
 def test_preprocess_dotproduct(tmpdir, fortran_reader, fortran_writer):