(Closes #2027) fix for ArrayMixin._effective_shape bug

examples/nemo/scripts/omp_cpu_trans.py

@@ -47,7 +47,6 @@
 
 # List of all files that psyclone will skip processing
 FILES_TO_SKIP = NOT_PERFORMANT + NOT_WORKING + [
-    "lbclnk.f90",  # TODO #2685: effective shape bug
     "asminc.f90",
     "trosk.f90",
     "vremap.f90",

examples/nemo/scripts/utils.py

@@ -37,8 +37,8 @@
 
 from psyclone.domain.common.transformations import KernelModuleInlineTrans
 from psyclone.psyir.nodes import (
-    Loop, Assignment, Directive, Container, Reference, CodeBlock, Call,
-    Return, IfBlock, Routine, IntrinsicCall)
+    Assignment, Loop, Directive, Container, Reference, CodeBlock,
+    Call, Return, IfBlock, Routine, IntrinsicCall)
 from psyclone.psyir.symbols import (
     DataSymbol, INTEGER_TYPE, REAL_TYPE, ArrayType, ScalarType,
     RoutineSymbol, ImportInterface)
@@ -246,6 +246,15 @@ def normalise_loops(
                     Reference2ArrayRangeTrans().apply(reference)
                 except TransformationError:
                     pass
+            if hasattr(reference, "indices"):
+                # Look at array-index expressions too.
+                for exprn in reference.indices:
+                    if (isinstance(exprn, Reference) and
+                            isinstance(exprn.symbol, DataSymbol)):
+                        try:
+                            Reference2ArrayRangeTrans().apply(exprn)
+                        except TransformationError:
+                            pass
 
     if loopify_array_intrinsics:
         for intr in schedule.walk(IntrinsicCall):

src/psyclone/psyir/frontend/fparser2.py

@@ -1337,6 +1337,7 @@ def _parse_dimensions(self, dimensions, symbol_table):
         shape = []
         # Traverse shape specs in Depth-first-search order
         for dim in walk(dimensions, (Fortran2003.Assumed_Shape_Spec,
+                                     Fortran2003.Deferred_Shape_Spec,
                                      Fortran2003.Explicit_Shape_Spec,
                                      Fortran2003.Assumed_Size_Spec)):
 
@@ -1367,6 +1368,11 @@ def _parse_dimensions(self, dimensions, symbol_table):
                 else:
                     shape.append(None)
 
+            elif isinstance(dim, Fortran2003.Deferred_Shape_Spec):
+                # Deferred_Shape_Spec has no children (R520). For our purposes
+                # it is equivalent to Assumed_Shape_Spec(None, None).
+                shape.append(None)
+
             elif isinstance(dim, Fortran2003.Explicit_Shape_Spec):
                 upper = self._process_array_bound(dim.items[1],
                                                   symbol_table)
@@ -1853,7 +1859,8 @@ def _process_decln(self, scope, symbol_table, decl, visibility_map=None,
         has_save_attr = False
         if attr_specs:
             for attr in attr_specs.items:
-                if isinstance(attr, Fortran2003.Attr_Spec):
+                if isinstance(attr, (Fortran2003.Attr_Spec,
+                                     Fortran2003.Component_Attr_Spec)):
                     normalized_string = str(attr).lower().replace(' ', '')
                     if normalized_string == "save":
                         if interface is not None:

src/psyclone/psyir/nodes/array_mixin.py

@@ -610,7 +610,8 @@ def _get_effective_shape(self):
         :rtype: list[:py:class:`psyclone.psyir.nodes.DataNode`]
 
         :raises NotImplementedError: if any of the array-indices involve a
-                                     function call or an expression.
+                                     function call or an expression or are
+                                     of unknown type.
         '''
         shape = []
         for idx, idx_expr in enumerate(self.indices):
@@ -621,7 +622,7 @@ def _get_effective_shape(self):
                 dtype = idx_expr.datatype
                 if isinstance(dtype, ArrayType):
                     # An array slice can be defined by a 1D slice of another
-                    # array, e.g. `a(b(1:4))`.
+                    # array, e.g. `a(b(1:4))` or `a(b)`.
                     indirect_array_shape = dtype.shape
                     if len(indirect_array_shape) > 1:
                         raise NotImplementedError(
@@ -630,8 +631,22 @@ def _get_effective_shape(self):
                             f"used to index into '{self.name}' has "
                             f"{len(indirect_array_shape)} dimensions.")
                     # pylint: disable=protected-access
-                    shape.append(idx_expr._extent(idx))
-
+                    if isinstance(idx_expr, ArrayMixin):
+                        shape.append(idx_expr._extent(idx))
+                    else:
+                        # We have a Reference (to an array) with no explicit
+                        # indexing. The extent of this is then the SIZE of
+                        # that array.
+                        sizeop = IntrinsicCall.create(
+                            IntrinsicCall.Intrinsic.SIZE, [idx_expr.copy()])
+                        shape.append(sizeop)
+                elif isinstance(dtype, (UnsupportedType, UnresolvedType)):
+                    raise NotImplementedError(
+                        f"The array index expression "
+                        f"'{idx_expr.debug_string()}' in access "
+                        f"'{self.debug_string()}' is of '{dtype}' type and "
+                        f"therefore whether it is an array slice (i.e. an "
+                        f"indirect access) cannot be determined.")
             elif isinstance(idx_expr, (Call, Operation, CodeBlock)):
                 # We can't yet straightforwardly query the type of a function
                 # call or Operation - TODO #1799.

src/psyclone/psyir/nodes/structure_reference.py

@@ -269,10 +269,34 @@ def datatype(self):
         :returns: the datatype of this reference.
         :rtype: :py:class:`psyclone.psyir.symbols.DataType`
 
-        :raises NotImplementedError: if the structure reference represents \
+        :raises NotImplementedError: if the structure reference represents
                                      an array of arrays.
 
         '''
+        def _get_cursor_shape(cursor, cursor_type):
+            '''
+            Utility that returns the shape of the supplied node.
+
+            :param cursor: the node to get the shape of.
+            :type cursor: :py:class:`psyclone.psyir.nodes.Node`
+            :param cursor_type: the type of the node.
+            :type cursor_type: :py:class:`psyclone.psyir.symbols.DataType`
+
+            :returns: the shape of the node or None if it is not an array.
+            :rtype: Optional[list[:py:class:`psyclone.psyir.nodes.Node`]]
+
+            '''
+            if not (isinstance(cursor, ArrayMixin) or
+                    isinstance(cursor_type, ArrayType)):
+                return None
+            if isinstance(cursor, ArrayMixin):
+                # It is an ArrayMixin and has indices so could be a single
+                # element or a slice.
+                # pylint: disable=protected-access
+                return cursor._get_effective_shape()
+            # No indices so it is an access to a whole array.
+            return cursor_type.shape
+
         # pylint: disable=too-many-return-statements, too-many-branches
         if self._overwrite_datatype:
             return self._overwrite_datatype
@@ -296,14 +320,10 @@ def datatype(self):
 
         # If the reference has explicit array syntax, we need to consider it
         # to calculate the resulting shape.
-        if isinstance(cursor, ArrayMixin):
-            try:
-                # pylint: disable=protected-access
-                shape = cursor._get_effective_shape()
-            except NotImplementedError:
-                return UnresolvedType()
-        else:
-            shape = []
+        try:
+            shape = _get_cursor_shape(cursor, cursor_type)
+        except NotImplementedError:
+            return UnresolvedType()
 
         # Walk down the structure, collecting information on any array slices
         # as we go.
@@ -317,37 +337,20 @@ def datatype(self):
                 # Once we've hit an Unresolved/UnsupportedType the cursor_type
                 # will remain set to that as we can't do any better.
                 cursor_type = cursor_type.components[cursor.name].datatype
-            if isinstance(cursor, ArrayMixin):
-                # pylint: disable=protected-access
-                try:
-                    shape.extend(cursor._get_effective_shape())
-                except NotImplementedError:
-                    return UnresolvedType()
-
-        # We've reached the ultimate member of the structure access.
-        if shape:
-            if isinstance(cursor_type, ArrayType):
-                # It's of array type but does it represent a single element,
-                # a slice or a whole array? (We use `children` rather than
-                # `indices` so as to avoid having to check that `cursor` is
-                # an `ArrayMember`.)
-                if cursor.children:
-                    # It has indices so could be a single element or a slice.
-                    # pylint: disable=protected-access
-                    cursor_shape = cursor._get_effective_shape()
-                else:
-                    # No indices so it is an access to a whole array.
-                    cursor_shape = cursor_type.shape
-                if cursor_shape and len(shape) != len(cursor_shape):
-                    # This ultimate access is an array but we've already
-                    # encountered one or more slices earlier in the access
-                    # expression.
+            try:
+                cursor_shape = _get_cursor_shape(cursor, cursor_type)
+            except NotImplementedError:
+                return UnresolvedType()
+            if cursor_shape:
+                if shape:
                     raise NotImplementedError(
                         f"Array of arrays not supported: the ultimate member "
                         f"'{cursor.name}' of the StructureAccess represents "
                         f"an array but other array notation is present in the "
                         f"full access expression: '{self.debug_string()}'")
+                shape = cursor_shape
 
+        if shape:
             return ArrayType(cursor_type, shape)
 
         # We don't have an explicit array access (because `shape` is Falsey)

src/psyclone/psyir/transformations/arrayassignment2loops_trans.py

@@ -49,10 +49,10 @@
     CodeBlock, Routine, BinaryOperation)
 from psyclone.psyir.nodes.array_mixin import ArrayMixin
 from psyclone.psyir.symbols import (
-    DataSymbol, INTEGER_TYPE, ScalarType, UnresolvedType, SymbolError,
-    ArrayType)
-from psyclone.psyir.transformations.transformation_error \
-    import TransformationError
+    ArrayType, DataSymbol, INTEGER_TYPE, ScalarType, SymbolError,
+    UnresolvedType)
+from psyclone.psyir.transformations.transformation_error import (
+    TransformationError)
 from psyclone.psyir.transformations.reference2arrayrange_trans import (
     Reference2ArrayRangeTrans)
 
@@ -241,17 +241,16 @@ def validate(self, node, options=None):
                 f"Range, but none were found in '{node.debug_string()}'."))
 
         # All the ArrayMixins must have the same number of Ranges to expand
-        found = None
+        num_of_ranges = None
         for accessor in node.walk(ArrayMixin):
-            num_of_ranges = len([x for x in accessor.indices
-                                 if isinstance(x, Range)])
-            if num_of_ranges > 0:
-                if not found:
+            count = len([x for x in accessor.indices if isinstance(x, Range)])
+            if count > 0:
+                if not num_of_ranges:
                     # If it's the first value we find, we store it
-                    found = num_of_ranges
+                    num_of_ranges = count
                 else:
                     # Otherwise we compare it against the previous found value
-                    if found != num_of_ranges:
+                    if count != num_of_ranges:
                         raise TransformationError(LazyString(
                             lambda: f"{self.name} does not support statements"
                             f" containing array accesses that have varying "
@@ -322,8 +321,8 @@ def validate(self, node, options=None):
                     lambda: f"{message} in:\n{node.debug_string()}."))
 
         # For each top-level reference (because we don't support nesting), the
-        # apply will have to be able to decide if its an Array (and access it
-        # with the index) or an Scalar (and leave it as it is). We can not
+        # apply() will have to determine whether it's an Array (and access it
+        # with the index) or a Scalar (and leave it as it is). We cannot
         # transform references where this is unclear.
         for reference in node.walk(Reference, stop_type=Reference):
             if isinstance(reference.parent, Call):
@@ -351,12 +350,12 @@ def validate(self, node, options=None):
             # scalar.
             if not isinstance(reference.datatype, (ScalarType, ArrayType)):
                 if isinstance(reference.symbol, DataSymbol):
-                    typestr = f"an {reference.symbol.datatype}"
+                    typestr = f"an {reference.datatype}"
                 else:
                     typestr = "not a DataSymbol"
                 message = (
                     f"{self.name} cannot expand expression because it "
-                    f"contains the variable '{reference.symbol.name}' "
+                    f"contains the access '{reference.debug_string()}' "
                     f"which is {typestr} and therefore cannot be guaranteed"
                     f" to be ScalarType.")
                 if not isinstance(reference.symbol, DataSymbol) or \

src/psyclone/tests/psyir/frontend/fparser2_where_handler_test.py

@@ -692,16 +692,24 @@ def test_where_stmt_validity():
 
 
 @pytest.mark.usefixtures("parser")
-def test_where_stmt():
+def test_where_stmt(fortran_reader):
     ''' Basic check that we handle a WHERE statement correctly.
 
     '''
-    fake_parent, _ = process_where(
-        "WHERE( at_i(:,:) > rn_amax_2d(:,:) )   "
-        "a_i(:,:,jl) = a_i(:,:,jl) * rn_amax_2d(:,:) / at_i(:,:)",
-        Fortran2003.Where_Stmt, ["at_i", "rn_amax_2d", "jl", "a_i"])
-    assert len(fake_parent.children) == 1
-    assert isinstance(fake_parent[0], Loop)
+    code = '''\
+    program where_test
+      implicit none
+      integer :: jl
+      real, dimension(:,:), allocatable :: at_i, rn_amax_2d
+      real, dimension(:,:,:), allocatable :: a_i
+      WHERE( at_i(:,:) > rn_amax_2d(:,:) ) &
+            a_i(:,:,jl) = a_i(:,:,jl) * rn_amax_2d(:,:) / at_i(:,:)
+    end program where_test
+    '''
+    psyir = fortran_reader.psyir_from_source(code)
+    routine = psyir.walk(Routine)[0]
+    assert len(routine.children) == 1
+    assert isinstance(routine[0], Loop)
 
 
 def test_where_stmt_no_reduction(fortran_reader, fortran_writer):
@@ -765,11 +773,11 @@ def test_where_ordering(parser):
      ("where (my_type%var(:) > epsi20)\n"
       "my_type2(jl)%array2(:,jl) = 3.0\n", "my_type%var"),
      ("where (my_type%block(jl)%var(:) > epsi20)\n"
-      "my_type%block%array(:,jl) = 3.0\n", "my_type%block(jl)%var"),
+      "my_type%block(jl)%array(:,jl) = 3.0\n", "my_type%block(jl)%var"),
      ("where (my_type%block(jl)%var(:) > epsi20)\n"
       "my_type%block(jl)%array(:,jl) = 3.0\n", "my_type%block(jl)%var"),
-     ("where (my_type2(:)%var > epsi20)\n"
-      "my_type2(:)%var = 3.0\n", "my_type2")])
+     ("where (my_type2(:)%var(jl) > epsi20)\n"
+      "my_type2(:)%var(jl) = 3.0\n", "my_type2")])
 def test_where_derived_type(fortran_reader, code, size_arg):
     ''' Test that we handle the case where array members of a derived type
     are accessed within a WHERE. '''
@@ -805,10 +813,12 @@ def test_where_derived_type(fortran_reader, code, size_arg):
     assert isinstance(loops[1].loop_body[0], IfBlock)
     # All Range nodes should have been replaced
     assert not loops[0].walk(Range)
-    # All ArrayMember accesses should now use the `widx1` loop variable
+    # All ArrayMember accesses other than 'var' should now use the `widx1`
+    # loop variable
     array_members = loops[0].walk(ArrayMember)
     for member in array_members:
-        assert "+ widx1 - 1" in member.indices[0].debug_string()
+        if member.name != "var":
+            assert "+ widx1 - 1" in member.indices[0].debug_string()
 
 
 @pytest.mark.parametrize(