diff --git a/src/psyclone/psyad/tl2ad.py b/src/psyclone/psyad/tl2ad.py index 040d308ef8..47133d66b8 100644 --- a/src/psyclone/psyad/tl2ad.py +++ b/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)) diff --git a/src/psyclone/psyad/transformations/preprocess.py b/src/psyclone/psyad/transformations/preprocess.py index a1051d2f25..9c9ef0b4d8 100644 --- a/src/psyclone/psyad/transformations/preprocess.py +++ b/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): diff --git a/src/psyclone/psyir/nodes/array_mixin.py b/src/psyclone/psyir/nodes/array_mixin.py index 7acb40380f..9a799794a3 100644 --- a/src/psyclone/psyir/nodes/array_mixin.py +++ b/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 diff --git a/src/psyclone/tests/psyad/tl2ad_test.py b/src/psyclone/tests/psyad/tl2ad_test.py index 843dd0e93c..710605efc0 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 3d4340f0d7..34d962ee73 100644 --- a/src/psyclone/tests/psyad/transformations/test_preprocess.py +++ b/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. diff --git a/src/psyclone/tests/psyir/nodes/array_mixin_test.py b/src/psyclone/tests/psyir/nodes/array_mixin_test.py index 7c4630fed1..dccacf19ce 100644 --- a/src/psyclone/tests/psyir/nodes/array_mixin_test.py +++ b/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