(Closes #2755) Bug fix for psyad failing to transform array assignment

src/psyclone/psyad/tl2ad.py

@@ -50,11 +50,14 @@
 from psyclone.psyad.transformations.preprocess import preprocess_trans
 from psyclone.psyir.backend.fortran import FortranWriter
 from psyclone.psyir.frontend.fortran import FortranReader
-from psyclone.psyir.nodes import Routine, Assignment, Reference, Literal, \
-    Call, Container, BinaryOperation, IntrinsicCall, ArrayReference, Range
-from psyclone.psyir.symbols import SymbolTable, ImportInterface, Symbol, \
-    ContainerSymbol, ScalarType, ArrayType, RoutineSymbol, DataSymbol, \
-    INTEGER_TYPE, UnresolvedType, UnsupportedType
+from psyclone.psyir.nodes import (
+    ArrayReference, Assignment, BinaryOperation, Call, Container,
+    IntrinsicCall, Literal, Range, Reference, Routine)
+from psyclone.psyir.symbols import (
+    SymbolTable, ImportInterface, Symbol,
+    ContainerSymbol, ScalarType, ArrayType, RoutineSymbol, DataSymbol,
+    INTEGER_TYPE, UnresolvedType, UnsupportedType)
+from psyclone.psyir.transformations import TransformationError
 
 
 #: The extent we will allocate to each dimension of arrays used in the
@@ -92,6 +95,7 @@ def generate_adjoint_str(tl_fortran_str, active_variables,
     :rtype: Tuple[str, str]
 
     :raises NotImplementedError: if the tangent-linear code is a function.
+    :raises NotImplementedError: if the pre-processing of the TL code fails.
     :raises NotImplementedError: if an unsupported API is specified.
 
     '''
@@ -115,7 +119,12 @@ def generate_adjoint_str(tl_fortran_str, active_variables,
 
     # Apply any required transformations to the TL PSyIR
     logger.debug("Preprocessing")
-    preprocess_trans(tl_psyir, active_variables)
+    try:
+        preprocess_trans(tl_psyir, active_variables)
+    except TransformationError as err:
+        raise NotImplementedError(
+            f"PSyAD failed to pre-process the supplied tangent-linear code. "
+            f"The error was: {str(err.value)}") from err
 
     logger.debug("PSyIR after TL preprocessing\n%s",
                  tl_psyir.view(colour=False))

src/psyclone/psyad/transformations/preprocess.py

@@ -42,7 +42,7 @@
 '''
 from psyclone.core import SymbolicMaths
 from psyclone.psyad.utils import node_is_passive
-from psyclone.psyir.nodes import (Assignment, IntrinsicCall, Reference)
+from psyclone.psyir.nodes import (Assignment, IntrinsicCall, Range, Reference)
 from psyclone.psyir.transformations import (DotProduct2CodeTrans,
                                             Matmul2CodeTrans,
                                             ArrayAssignment2LoopsTrans,
@@ -58,16 +58,19 @@ def preprocess_trans(kernel_psyir, active_variable_names):
     called internally by the PSyAD script before transforming the code
     to its adjoint form.
 
-    :param kernel_psyir: PSyIR representation of the tangent linear \
+    :param kernel_psyir: PSyIR representation of the tangent linear
         kernel code.
     :type kernel_psyir: :py:class:`psyclone.psyir.nodes.Node`
     :param active_variable_names: list of active variable names.
-    :type active_variable_names: list of str
+    :type active_variable_names: list[str]
+
+    :raises TransformationError: if an active array assignment cannot be
+        transformed into an explicit loop.
 
     '''
     dot_product_trans = DotProduct2CodeTrans()
     matmul_trans = Matmul2CodeTrans()
-    arrayrange2loop_trans = ArrayAssignment2LoopsTrans()
+    arrayassign2loops_trans = ArrayAssignment2LoopsTrans()
     reference2arrayrange_trans = Reference2ArrayRangeTrans()
 
     # Replace references to arrays (array notation) with array-ranges
@@ -77,19 +80,6 @@ def preprocess_trans(kernel_psyir, active_variable_names):
         except TransformationError:
             pass
 
-    # Replace array-ranges with explicit loops
-    for assignment in kernel_psyir.walk(Assignment):
-        if node_is_passive(assignment, active_variable_names):
-            # No need to modify passive assignments
-            continue
-        # Repeatedly apply the transformation until there are no more
-        # array ranges in this assignment.
-        while True:
-            try:
-                arrayrange2loop_trans.apply(assignment)
-            except TransformationError:
-                break
-
     for call in kernel_psyir.walk(IntrinsicCall):
         if call.intrinsic == IntrinsicCall.Intrinsic.DOT_PRODUCT:
             # Apply DOT_PRODUCT transformation
@@ -98,6 +88,20 @@ def preprocess_trans(kernel_psyir, active_variable_names):
             # Apply MATMUL transformation
             matmul_trans.apply(call)
 
+    # Replace array-ranges with explicit loops
+    for assignment in kernel_psyir.walk(Assignment):
+        if node_is_passive(assignment, active_variable_names):
+            # No need to modify passive assignments
+            continue
+        try:
+            arrayassign2loops_trans.apply(assignment)
+        except TransformationError:
+            # Double-check that the transformation succeeded in
+            # handling any explicit array assignments. If it didn't then
+            # we can't create the adjoint.
+            if assignment.lhs.walk(Range):
+                raise
+
     # Deal with any associativity issues here as AssignmentTrans
     # is not able to.
     for assignment in kernel_psyir.walk(Assignment):

src/psyclone/psyir/nodes/array_mixin.py

@@ -610,15 +610,14 @@ 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 or are
-                                     of unknown type.
+                                     function call or are of unknown type.
         '''
         shape = []
         for idx, idx_expr in enumerate(self.indices):
             if isinstance(idx_expr, Range):
                 shape.append(self._extent(idx))
 
-            elif isinstance(idx_expr, Reference):
+            elif isinstance(idx_expr, (Reference, Operation)):
                 dtype = idx_expr.datatype
                 if isinstance(dtype, ArrayType):
                     # An array slice can be defined by a 1D slice of another
@@ -634,9 +633,9 @@ def _get_effective_shape(self):
                     if isinstance(idx_expr, ArrayMixin):
                         shape.append(idx_expr._extent(idx))
                     else:
-                        # We have a Reference (to an array) with no explicit
+                        # We have some expression with a shape but no explicit
                         # indexing. The extent of this is then the SIZE of
-                        # that array.
+                        # that array (expression).
                         sizeop = IntrinsicCall.create(
                             IntrinsicCall.Intrinsic.SIZE, [idx_expr.copy()])
                         shape.append(sizeop)
@@ -647,13 +646,13 @@ def _get_effective_shape(self):
                         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)):
+            elif isinstance(idx_expr, (Call, CodeBlock)):
                 # We can't yet straightforwardly query the type of a function
-                # call or Operation - TODO #1799.
+                # call - TODO #1799.
                 raise NotImplementedError(
                     f"The array index expressions for access "
                     f"'{self.debug_string()}' include a function call or "
-                    f"expression. Querying the return type of "
+                    f"unsupported feature. Querying the return type of "
                     f"such things is yet to be implemented.")
 
         return shape

src/psyclone/tests/psyad/tl2ad_test.py

@@ -207,6 +207,26 @@ def test_generate_adjoint_str_trans(tmpdir):
     assert Compile(tmpdir).string_compiles(result)
 
 
+def test_generate_adjoint_str_trans_error(tmpdir):
+    '''Test that the generate_adjoint_str() function successfully catches
+    an error from the preprocess_trans() function.
+
+    '''
+    code = (
+        "program test\n"
+        "use other_mod, only: func\n"
+        "real, dimension(10,10,10) :: a,b,c,d,e,f\n"
+        "integer, dimension(10) :: map\n"
+        "integer, parameter :: i = 5\n"
+        "a(:,1,:) = b(:,1,:) * c(:,1+int(real(complex(1.0,1.0))),:)\n"
+        "end program test\n")
+    with pytest.raises(NotImplementedError) as err:
+        _ = generate_adjoint_str(code, ["a", "c"])
+    assert ("failed to pre-process the supplied tangent-linear code. The error"
+            " was: Transformation Error: ArrayAssignment2LoopsTrans does not"
+            in str(err.value))
+
+
 def test_generate_adjoint_str_generate_harness_no_api(tmpdir):
     '''Test the create_test option to generate_adjoint_str() when no
     API is specified.'''

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

@@ -43,6 +43,7 @@
 from psyclone.psyad.transformations.preprocess import preprocess_trans
 from psyclone.psyir.backend.fortran import FortranWriter
 from psyclone.psyir.frontend.fortran import FortranReader
+from psyclone.psyir.transformations import TransformationError
 from psyclone.tests.utilities import Compile
 
 
@@ -182,7 +183,7 @@ def test_preprocess_matmul(tmpdir, fortran_reader, fortran_writer):
     assert Compile(tmpdir).string_compiles(result)
 
 
-def test_preprocess_arrayrange2loop(tmpdir, fortran_reader, fortran_writer):
+def test_preprocess_arrayassign2loop(tmpdir, fortran_reader, fortran_writer):
     '''Test that the preprocess script replaces active assignments that
     contain arrays that use range notation with equivalent code that
     uses explicit loops. Also check that they are not modified if they
@@ -192,24 +193,19 @@ def test_preprocess_arrayrange2loop(tmpdir, fortran_reader, fortran_writer):
     code = (
         "program test\n"
         "real, dimension(10,10,10) :: a,b,c,d,e,f\n"
-        "a(:,1,:) = b(:,1,:) * c(:,1,:)\n"
+        "integer, dimension(10) :: map\n"
+        "integer, parameter :: i = 5\n"
+        "a(:,1,:) = b(:,1,:) * c(:,1+map(i),:)\n"
         "d(1,1,1) = 0.0\n"
         "e(:,:,:) = f(:,:,:)\n"
         "print *, \"hello\"\n"
         "end program test\n")
     expected = (
-        "program test\n"
-        "  real, dimension(10,10,10) :: a\n"
-        "  real, dimension(10,10,10) :: b\n"
-        "  real, dimension(10,10,10) :: c\n"
-        "  real, dimension(10,10,10) :: d\n"
-        "  real, dimension(10,10,10) :: e\n"
-        "  real, dimension(10,10,10) :: f\n"
         "  integer :: idx\n"
         "  integer :: idx_1\n\n"
         "  do idx = LBOUND(a, dim=3), UBOUND(a, dim=3), 1\n"
         "    do idx_1 = LBOUND(a, dim=1), UBOUND(a, dim=1), 1\n"
-        "      a(idx_1,1,idx) = b(idx_1,1,idx) * c(idx_1,1,idx)\n"
+        "      a(idx_1,1,idx) = b(idx_1,1,idx) * c(idx_1,map(i) + 1,idx)\n"
         "    enddo\n"
         "  enddo\n"
         "  d(1,1,1) = 0.0\n"
@@ -221,10 +217,31 @@ def test_preprocess_arrayrange2loop(tmpdir, fortran_reader, fortran_writer):
     psyir = fortran_reader.psyir_from_source(code)
     preprocess_trans(psyir, ["a", "c"])
     result = fortran_writer(psyir)
-    assert result == expected
+    assert expected in result
     assert Compile(tmpdir).string_compiles(result)
 
 
+def test_preprocess_arrayassign2loop_failure(fortran_reader, fortran_writer):
+    '''
+    Check that we catch the case where the array-assignment transformation
+    fails.
+
+    '''
+    code = (
+        "program test\n"
+        "use other_mod, only: func\n"
+        "real, dimension(10,10,10) :: a,b,c,d,e,f\n"
+        "integer, dimension(10) :: map\n"
+        "integer, parameter :: i = 5\n"
+        "a(:,1,:) = b(:,1,:) * c(:,1+int(real(complex(1.0,1.0))),:)\n"
+        "end program test\n")
+    psyir = fortran_reader.psyir_from_source(code)
+    with pytest.raises(TransformationError) as err:
+        preprocess_trans(psyir, ["a", "c"])
+    assert (" ArrayAssignment2LoopsTrans does not accept calls which are "
+            "not guaranteed" in str(err.value))
+
+
 @pytest.mark.parametrize("operation", ["+", "-"])
 def test_preprocess_associativity(operation, fortran_reader, fortran_writer):
     '''Test that associativity is handled correctly.

src/psyclone/tests/psyir/nodes/array_mixin_test.py

@@ -622,6 +622,7 @@ def test_get_effective_shape(fortran_reader):
         "  b(indices(2:3,1:2), 2:5) = 2.0\n"
         "  a(f()) = 2.0\n"
         "  a(2+3) = 1.0\n"
+        "  b(idx, 1+indices(1,1):) = 1\n"
         "  b(idx, a) = -1.0\n"
         "  b(scalarval, arrayval) = 1\n"
         "end subroutine\n")
@@ -688,12 +689,17 @@ def test_get_effective_shape(fortran_reader):
     child_idx += 1
     with pytest.raises(NotImplementedError) as err:
         _ = routine.children[child_idx].lhs._get_effective_shape()
-    assert "include a function call or expression" in str(err.value)
-    # Array access with expression in indices.
+    assert "include a function call or unsupported feature" in str(err.value)
+    # Array access with simple expression in indices.
     child_idx += 1
-    with pytest.raises(NotImplementedError) as err:
-        _ = routine.children[child_idx].lhs._get_effective_shape()
-    assert "include a function call or expression" in str(err.value)
+    shape = routine.children[child_idx].lhs._get_effective_shape()
+    assert shape == []
+    # Array access with expression involving indirect access in indices.
+    child_idx += 1
+    shape = routine.children[child_idx].lhs._get_effective_shape()
+    assert len(shape) == 1
+    assert (shape[0].debug_string().lower() ==
+            "ubound(b, dim=2) - (1 + indices(1,1)) + 1")
     # Array access with indices given by another array that is not explicitly
     # indexed.
     child_idx += 1