Repository: stfc/PSyclone (python)

Issue:
[LFRic] [PSyAD] PSyAD no longer computes adjoint of vector assigments properly in PSyclone 3.0-dev
In testing PSyclone 3.0, we find that forward code that looks like

```
vector(:) = vector(:) + constant(:) * scalar
```

where `vector` and `scalar` are active no longer adjoints to

```
do idx = vector_ubound, vector_lbound, -1
  scalar = scalar + constant(idx) * vector(idx)
end do
```

but instead does

```
scalar = scalar + constant(:) * vector(:)
```

which fails to compile. The former behaviour is expected.

Pull request description:
(Closes #2755) Bug fix for psyad failing to transform array assignment
A small change that (pleasingly) only involved me updating some code to remove an old restriction. @DrTVockerodtMO please could you try this branch and let me know whether it solves your problem?

Patch:
diff --git a/changelog b/changelog
index b806b6b7c..551dd689d 100644
--- a/changelog
+++ b/changelog
@@ -281,6 +281,9 @@
 
 	107) PR #2776 for #1280. Fixes ref_guide warnings.
 
+	108) PR #2767 for #2755. Bug fix for psyad failing to transform array
+	assignment
+
 release 2.5.0 14th of February 2024
 
 	1) PR #2199 for #2189. Fix bugs with missing maps in enter data
diff --git a/doc/developer_guide/index.rst b/doc/developer_guide/index.rst
index 686feed65..3fb52ad3d 100644
--- a/doc/developer_guide/index.rst
+++ b/doc/developer_guide/index.rst
@@ -57,7 +57,7 @@
    psy_data
    system_specific_setup
    coding-style
-   zz_bibliography
+   Bibliography <zz_bibliography>
 
 ..   FAQS
 
diff --git a/doc/developer_guide/system_specific_setup.rst b/doc/developer_guide/system_specific_setup.rst
index c76360f8a..41d2b55fe 100644
--- a/doc/developer_guide/system_specific_setup.rst
+++ b/doc/developer_guide/system_specific_setup.rst
@@ -3,7 +3,7 @@
 System-specific Developer Set-up
 ================================
 
-Section :ref:`user_guide:system_specific_setup` in the PSyclone User Guide
+Section :ref:`user_guide:getting-going-download` in the PSyclone User Guide
 describes the setup for a user of PSyclone. It includes all steps necessary
 to be able to use PSyclone. And while you could obviously do
 some development, none of the required tools for testing or
@@ -12,7 +12,7 @@ documentation creation will be installed.
 This section adds software that is used to develop and test
 PSyclone. It includes all packages for testing and creation of
 documentation in html and pdf. We assume you have already installed
-the software described in the :ref:`user_guide:system_specific_setup` section.
+the software described in the :ref:`user_guide:getting-going-download` section.
 
 It contains detailed instructions for Ubuntu 16.04.2 and 
 OpenSUSE 42.2 - if you are working with a different Linux
diff --git a/src/psyclone/psyad/tl2ad.py b/src/psyclone/psyad/tl2ad.py
index 040d308ef..47133d66b 100644
--- a/src/psyclone/psyad/tl2ad.py
+++ b/src/psyclone/psyad/tl2ad.py
@@ -50,11 +50,14 @@ from psyclone.psyad.domain.lfric import (generate_lfric_adjoint,
 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))
diff --git a/src/psyclone/psyad/transformations/preprocess.py b/src/psyclone/psyad/transformations/preprocess.py
index a1051d2f2..adfc09092 100644
--- a/src/psyclone/psyad/transformations/preprocess.py
+++ b/src/psyclone/psyad/transformations/preprocess.py
@@ -42,7 +42,7 @@ translated to adjoint PSyIR.
 '''
 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,16 @@ 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
+        # Earlier use of Reference2ArrayRangeTrans will ensure LHS has a Range
+        # if it is an array assignment.
+        if assignment.lhs.walk(Range):
+            arrayassign2loops_trans.apply(assignment)
+
     # Deal with any associativity issues here as AssignmentTrans
     # is not able to.
     for assignment in kernel_psyir.walk(Assignment):
diff --git a/src/psyclone/psyir/nodes/array_mixin.py b/src/psyclone/psyir/nodes/array_mixin.py
index 7acb40380..9a799794a 100644
--- a/src/psyclone/psyir/nodes/array_mixin.py
+++ b/src/psyclone/psyir/nodes/array_mixin.py
@@ -610,15 +610,14 @@ class ArrayMixin(metaclass=abc.ABCMeta):
         :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 @@ class ArrayMixin(metaclass=abc.ABCMeta):
                     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 @@ class ArrayMixin(metaclass=abc.ABCMeta):
                         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


Test patch:
diff --git a/src/psyclone/tests/psyad/tl2ad_test.py b/src/psyclone/tests/psyad/tl2ad_test.py
index 843dd0e93..710605efc 100644
--- a/src/psyclone/tests/psyad/tl2ad_test.py
+++ b/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.'''
diff --git a/src/psyclone/tests/psyad/transformations/test_preprocess.py b/src/psyclone/tests/psyad/transformations/test_preprocess.py
index 3d4340f0d..34d962ee7 100644
--- a/src/psyclone/tests/psyad/transformations/test_preprocess.py
+++ b/src/psyclone/tests/psyad/transformations/test_preprocess.py
@@ -43,6 +43,7 @@ import pytest
 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.
diff --git a/src/psyclone/tests/psyir/nodes/array_mixin_test.py b/src/psyclone/tests/psyir/nodes/array_mixin_test.py
index 7c4630fed..b3915217c 100644
--- a/src/psyclone/tests/psyir/nodes/array_mixin_test.py
+++ b/src/psyclone/tests/psyir/nodes/array_mixin_test.py
@@ -622,38 +622,45 @@ 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")
     psyir = fortran_reader.psyir_from_source(code)
     routine = psyir.walk(Routine)[0]
     # Direct array slice.
+    #   a(1:10) = 0.0
     child_idx = 0
     shape = routine.children[child_idx].lhs._get_effective_shape()
     assert len(shape) == 1
     assert isinstance(shape[0], Literal)
     assert shape[0].value == "10"
     # Array slice with non-unit step.
+    #   a(1:10:3) = 0.0
     child_idx += 1
     shape = routine.children[child_idx].lhs._get_effective_shape()
     assert len(shape) == 1
     assert shape[0].debug_string() == "(10 - 1) / 3 + 1"
     # Full array slice without bounds.
+    #   a(:) = 0.0
     child_idx += 1
     shape = routine.children[child_idx].lhs._get_effective_shape()
     assert len(shape) == 1
     assert "SIZE(a, dim=1)" in shape[0].debug_string()
     # Array slice with only lower-bound specified.
+    #   a(2:) = 0.0
     child_idx += 1
     shape = routine.children[child_idx].lhs._get_effective_shape()
     assert len(shape) == 1
     assert shape[0].debug_string() == "UBOUND(a, dim=1) - 2 + 1"
     # Array slice with only upper-bound specified.
+    #   a(:5) = 0.0
     child_idx += 1
     shape = routine.children[child_idx].lhs._get_effective_shape()
     assert len(shape) == 1
     assert shape[0].debug_string() == "5 - LBOUND(a, dim=1) + 1"
     # Array slice with only step specified.
+    #   a(::4) = 0.0
     child_idx += 1
     shape = routine.children[child_idx].lhs._get_effective_shape()
     assert len(shape) == 1
@@ -663,11 +670,13 @@ def test_get_effective_shape(fortran_reader):
             "(UBOUND(a, dim=1) - LBOUND(a, dim=1)) / 4 + 1")
     # Array slice defined using LBOUND and UBOUND intrinsics but for a
     # different array altogether.
+    #   a(lbound(b,1):ubound(b,2)) = 0.0
     child_idx += 1
     shape = routine.children[child_idx].lhs._get_effective_shape()
     assert len(shape) == 1
     assert shape[0].debug_string() == "UBOUND(b, 2) - LBOUND(b, 1) + 1"
     # Indirect array slice.
+    #   b(indices(2:3,1), 2:5) = 2.0
     child_idx += 1
     shape = routine.children[child_idx].lhs._get_effective_shape()
     assert len(shape) == 2
@@ -679,29 +688,40 @@ def test_get_effective_shape(fortran_reader):
     assert shape[1].debug_string() == "5 - 2 + 1"
     # An indirect array slice can only be 1D.
     child_idx += 1
+    #   b(indices(2:3,1:2), 2:5) = 2.0
     with pytest.raises(NotImplementedError) as err:
         _ = routine.children[child_idx].lhs._get_effective_shape()
     assert ("array defining a slice of a dimension of another array must be "
             "1D but 'indices' used to index into 'b' has 2 dimensions" in
             str(err.value))
     # Indirect array access using function call.
+    #   a(f()) = 2.0
     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.
+    #   a(2+3) = 1.0
     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.
+    #   b(idx, 1+indices(1,1):) = 1
+    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.
+    #   b(idx, a) = -1.0
     child_idx += 1
     shape = routine.children[child_idx].lhs._get_effective_shape()
     assert len(shape) == 1
     assert "SIZE(a)" in shape[0].debug_string()
     # Array-index expressions are symbols of unknown type so we don't know
     # whether we have an array slice or just a scalar.
+    #   b(scalarval, arrayval) = 1
     child_idx += 1
     with pytest.raises(NotImplementedError) as err:
         _ = routine.children[child_idx].lhs._get_effective_shape()