Repository: stfc/PSyclone

Issue:
[psyir] Assignment.is_array_range does not work for Structures
Currently, `Assignment.is_array_range` only checks that the lhs is an `ArrayReference`. Now that we can have e.g.

    some_struc(ji)%data(:,:) = 0.0

this is clearly insufficient and results in missed parallelisation opportunities when processing NEMO, as reported by @deardenchris .

Discussion:

Patch:
diff --git a/changelog b/changelog
index d89377fea..9338cac02 100644
--- a/changelog
+++ b/changelog
@@ -453,6 +453,9 @@
         that the 'same space' option is now provided with the options
         dict, in line with the way other transformations work.
 
+	147) PR #1165 for #1164. Update PSyIR Assignment node
+	is_array_range test to work with structures.
+
 release 1.9.0 20th May 2020
 
 	1) #602 for #597. Modify Node.ancestor() to optionally include
diff --git a/src/psyclone/psyir/nodes/assignment.py b/src/psyclone/psyir/nodes/assignment.py
index b4d6772b3..1a0a99d0c 100644
--- a/src/psyclone/psyir/nodes/assignment.py
+++ b/src/psyclone/psyir/nodes/assignment.py
@@ -39,9 +39,12 @@
 ''' This module contains the Assignment node implementation.'''
 
 import re
+from psyclone.psyir.nodes.array_reference import ArrayReference
+from psyclone.psyir.nodes.ranges import Range
 from psyclone.psyir.nodes.statement import Statement
 from psyclone.psyir.nodes.datanode import DataNode
 from psyclone.psyir.nodes.structure_reference import StructureReference
+from psyclone.psyir.nodes.operation import Operation, REDUCTION_OPERATORS
 from psyclone.core.access_info import VariablesAccessInfo, AccessType
 from psyclone.errors import InternalError
 from psyclone.f2pygen import PSyIRGen
@@ -188,16 +191,37 @@ class Assignment(Statement):
     @property
     def is_array_range(self):
         '''
-        returns: True if the lhs of the assignment is an array with at \
-            least one of its dimensions being a range and False \
-            otherwise.
+        returns: True if the lhs of the assignment is an array access with at \
+            least one of its dimensions being a range and False otherwise.
         rtype: bool
 
         '''
-        from psyclone.psyir.nodes import ArrayReference, Range
-        if not isinstance(self.lhs, ArrayReference):
-            return False
-        return any(dim for dim in self.lhs.children if isinstance(dim, Range))
+        # It's not sufficient simply to check for a Range node as that may be
+        # part of an argument to an Operator or function that performs a
+        # reduction and thus returns a scalar result, e.g. a(SUM(b(:))) = 1.0
+        # TODO #658 this check for reductions needs extending to also support
+        # user-implemented functions.
+        if isinstance(self.lhs, (ArrayReference, StructureReference)):
+            ranges = self.lhs.walk(Range)
+            for array_range in ranges:
+                opn = array_range.ancestor(Operation)
+                while opn:
+                    if opn.operator in REDUCTION_OPERATORS:
+                        # The current array range is in an argument to a
+                        # reduction operation so we assume that the result
+                        # is a scalar.
+                        # TODO 658 this could still be a reduction into an
+                        # array (e.g. SUM(a(:,:), 1)) but we need to be able
+                        # to interrogate the type of a PSyIR expression in
+                        # order to be sure. e.g. SUM(a(:,:), mask(:,:)) will
+                        # return a scalar.
+                        break
+                    opn = opn.ancestor(Operation)
+                else:
+                    # We didn't find a reduction operation so there is an
+                    # array range on the LHS
+                    return True
+        return False
 
     def gen_code(self, parent):
         '''F2pygen code generation of an Assignment.
diff --git a/src/psyclone/psyir/nodes/operation.py b/src/psyclone/psyir/nodes/operation.py
index 6b607e1b2..e5102c257 100644
--- a/src/psyclone/psyir/nodes/operation.py
+++ b/src/psyclone/psyir/nodes/operation.py
@@ -398,3 +398,14 @@ class NaryOperation(Operation):
         for child in children:
             child.parent = nary_op
         return nary_op
+
+
+# TODO #658 this can be removed once we have support for determining the
+# type of a PSyIR expression.
+#: Those operators that perform a reduction on an array.
+REDUCTION_OPERATORS = [UnaryOperation.Operator.SUM,
+                       BinaryOperation.Operator.SUM,
+                       NaryOperation.Operator.SUM]
+
+# For automatic API documentation generation
+__all__ = ["Operation", "UnaryOperation", "BinaryOperation", "NaryOperation"]


Test patch:
diff --git a/src/psyclone/tests/psyir/nodes/assignment_test.py b/src/psyclone/tests/psyir/nodes/assignment_test.py
index ac1ad34dc..a753df704 100644
--- a/src/psyclone/tests/psyir/nodes/assignment_test.py
+++ b/src/psyclone/tests/psyir/nodes/assignment_test.py
@@ -41,9 +41,11 @@
 from __future__ import absolute_import
 import pytest
 from psyclone.psyir.nodes import Assignment, Reference, Literal, \
-    ArrayReference, Range
-from psyclone.psyir.symbols import DataSymbol, REAL_SINGLE_TYPE, \
-    INTEGER_SINGLE_TYPE, REAL_TYPE, ArrayType, INTEGER_TYPE
+    ArrayReference, Range, BinaryOperation, StructureReference, \
+    ArrayOfStructuresReference, UnaryOperation
+from psyclone.psyir.symbols import DataSymbol, REAL_SINGLE_TYPE, Symbol, \
+    INTEGER_SINGLE_TYPE, REAL_TYPE, ArrayType, INTEGER_TYPE, StructureType, \
+    TypeSymbol
 from psyclone.errors import InternalError, GenerationError
 from psyclone.psyir.backend.fortran import FortranWriter
 from psyclone.tests.utilities import check_links
@@ -133,28 +135,182 @@ def test_assignment_children_validation():
 
 def test_is_array_range():
     '''test that the is_array_range method behaves as expected, returning
-    true if the LHS of the assignment is an array range access and
-    false otherwise.
+    true if the LHS of the assignment is an array range access.
 
     '''
     one = Literal("1.0", REAL_TYPE)
     int_one = Literal("1", INTEGER_TYPE)
+    int_ten = Literal("10", INTEGER_TYPE)
+
+    # lhs is an array reference with a range
+    array_type = ArrayType(REAL_TYPE, [10, 10])
+    symbol = DataSymbol("x", array_type)
+    x_range = Range.create(int_one, int_ten.copy(), int_one.copy())
+    array_ref = ArrayReference.create(symbol, [x_range, int_one.copy()])
+    assignment = Assignment.create(array_ref, one.copy())
+    assert assignment.is_array_range is True
+
+    # Check when lhs consists of various forms of structure access
+    grid_type = StructureType.create([
+        ("dx", REAL_SINGLE_TYPE, Symbol.Visibility.PUBLIC),
+        ("dy", REAL_SINGLE_TYPE, Symbol.Visibility.PUBLIC)])
+    grid_type_symbol = TypeSymbol("grid_type", grid_type)
+    # Create the definition of the 'field_type', contains array of grid_types
+    field_type_def = StructureType.create(
+        [("data", ArrayType(REAL_SINGLE_TYPE, [10]), Symbol.Visibility.PUBLIC),
+         ("sub_meshes", ArrayType(grid_type_symbol, [3]),
+          Symbol.Visibility.PUBLIC)])
+    field_type_symbol = TypeSymbol("field_type", field_type_def)
+    field_symbol = DataSymbol("wind", field_type_symbol)
+
+    # Array reference to component of derived type using a range
+    lbound = BinaryOperation.create(
+        BinaryOperation.Operator.LBOUND,
+        StructureReference.create(field_symbol, ["data"]), int_one.copy())
+    ubound = BinaryOperation.create(
+        BinaryOperation.Operator.UBOUND,
+        StructureReference.create(field_symbol, ["data"]), int_one.copy())
+    my_range = Range.create(lbound, ubound)
+
+    data_ref = StructureReference.create(field_symbol, [("data", [my_range])])
+    assign = Assignment.create(data_ref, one.copy())
+    assert assign.is_array_range is True
+
+    # Access to slice of 'sub_meshes': wind%sub_meshes(1:3)%dx = 1.0
+    sub_range = Range.create(int_one.copy(), Literal("3", INTEGER_TYPE))
+    dx_ref = StructureReference.create(field_symbol, [("sub_meshes",
+                                                       [sub_range]), "dx"])
+    sub_assign = Assignment.create(dx_ref, one.copy())
+    assert sub_assign.is_array_range is True
+
+    # Create an array of these derived types and assign to a slice:
+    # chi(1:10)%data(1) = 1.0
+    field_bundle_symbol = DataSymbol("chi", ArrayType(field_type_symbol, [3]))
+    fld_range = Range.create(int_one.copy(), Literal("10", INTEGER_TYPE))
+    fld_ref = ArrayOfStructuresReference.create(field_bundle_symbol,
+                                                [fld_range],
+                                                [("data", [int_one.copy()])])
+    fld_assign = Assignment.create(fld_ref, one.copy())
+    assert fld_assign.is_array_range is True
+
+    # When the slice has two operator ancestors, none of which are a reduction
+    # e.g y(1, INT(ABS(map(:, 1)))) = 1.0
+    int_array_type = ArrayType(INTEGER_SINGLE_TYPE, [10, 10])
+    map_sym = DataSymbol("map", int_array_type)
+    lbound1 = BinaryOperation.create(
+        BinaryOperation.Operator.LBOUND,
+        Reference(map_sym), int_one.copy())
+    ubound1 = BinaryOperation.create(
+        BinaryOperation.Operator.UBOUND,
+        Reference(map_sym), int_one.copy())
+    my_range1 = Range.create(lbound1, ubound1)
+    abs_op = UnaryOperation.create(UnaryOperation.Operator.ABS,
+                                   ArrayReference.create(map_sym,
+                                                         [my_range1,
+                                                          int_one.copy()]))
+    int_op = UnaryOperation.create(UnaryOperation.Operator.INT, abs_op)
+    assignment = Assignment.create(
+        ArrayReference.create(symbol, [int_one.copy(), int_op]),
+        one.copy())
+    assert assignment.is_array_range is True
+
+
+@pytest.mark.xfail(reason="#658 needs typing of PSyIR expressions")
+def test_array_range_with_reduction():
+    ''' Test that we correctly identify an array range when it is the result
+        of a reduction from an array, e.g x(1, INT(SUM(map(:, :), 1))) = 1.0
+
+    '''
+    one = Literal("1.0", REAL_TYPE)
+    int_one = Literal("1", INTEGER_TYPE)
+    int_two = Literal("2", INTEGER_TYPE)
+    int_array_type = ArrayType(INTEGER_SINGLE_TYPE, [10, 10])
+    map_sym = DataSymbol("map", int_array_type)
+    array_type = ArrayType(REAL_TYPE, [10, 10])
+    symbol = DataSymbol("x", array_type)
+    lbound1 = BinaryOperation.create(
+        BinaryOperation.Operator.LBOUND,
+        Reference(map_sym), int_one.copy())
+    ubound1 = BinaryOperation.create(
+        BinaryOperation.Operator.UBOUND,
+        Reference(map_sym), int_one.copy())
+    my_range1 = Range.create(lbound1, ubound1)
+    lbound2 = BinaryOperation.create(
+        BinaryOperation.Operator.LBOUND,
+        Reference(map_sym), int_two.copy())
+    ubound2 = BinaryOperation.create(
+        BinaryOperation.Operator.UBOUND,
+        Reference(map_sym), int_two.copy())
+    my_range2 = Range.create(lbound2, ubound2)
+    bsum_op = BinaryOperation.create(BinaryOperation.Operator.SUM,
+                                     ArrayReference.create(map_sym,
+                                                           [my_range1,
+                                                            my_range2]),
+                                     int_one.copy())
+    int_op2 = UnaryOperation.create(UnaryOperation.Operator.INT, bsum_op)
+    assignment = Assignment.create(
+        ArrayReference.create(symbol,
+                              [int_one.copy(), int_op2]),
+        one.copy())
+    assert assignment.is_array_range is True
+
+
+def test_is_not_array_range():
+    ''' Test that is_array_range correctly rejects things that aren't
+    an assignment to an array range.
+
+    '''
+    int_one = Literal("1", INTEGER_SINGLE_TYPE)
+    one = Literal("1.0", REAL_TYPE)
     var = DataSymbol("x", REAL_TYPE)
     reference = Reference(var)
 
     # lhs is not an array
     assignment = Assignment.create(reference, one)
-    assert not assignment.is_array_range
+    assert assignment.is_array_range is False
 
     # lhs is an array reference but has no range
     array_type = ArrayType(REAL_TYPE, [10, 10])
-    symbol = DataSymbol("x", array_type)
+    symbol = DataSymbol("y", array_type)
     array_ref = Reference(symbol)
     assignment = Assignment.create(array_ref, one.copy())
-    assert not assignment.is_array_range
+    assert assignment.is_array_range is False
 
-    # lhs is an array reference with a range
-    my_range = Range.create(int_one, int_one.copy(), int_one.copy())
-    array_ref = ArrayReference.create(symbol, [my_range, int_one.copy()])
-    assignment = Assignment.create(array_ref, one.copy())
-    assert assignment.is_array_range
+    # lhs is an array reference but the single index value is obtained
+    # using an array range, y(1, SUM(map(:), 1)) = 1.0
+    int_array_type = ArrayType(INTEGER_SINGLE_TYPE, [10])
+    map_sym = DataSymbol("map", int_array_type)
+    start = BinaryOperation.create(BinaryOperation.Operator.LBOUND,
+                                   Reference(map_sym), int_one.copy())
+    stop = BinaryOperation.create(BinaryOperation.Operator.UBOUND,
+                                  Reference(map_sym), int_one.copy())
+    my_range = Range.create(start, stop)
+    sum_op = BinaryOperation.create(BinaryOperation.Operator.SUM,
+                                    ArrayReference.create(map_sym, [my_range]),
+                                    int_one.copy())
+    assignment = Assignment.create(
+        ArrayReference.create(symbol, [int_one.copy(), sum_op]),
+        one.copy())
+    assert assignment.is_array_range is False
+
+    # When the slice has two operator ancestors, one of which is a reduction
+    # e.g y(1, SUM(ABS(map(:)), 1)) = 1.0
+    abs_op = UnaryOperation.create(UnaryOperation.Operator.ABS,
+                                   ArrayReference.create(map_sym,
+                                                         [my_range.copy()]))
+    sum_op2 = BinaryOperation.create(BinaryOperation.Operator.SUM,
+                                     abs_op, int_one.copy())
+    assignment = Assignment.create(
+        ArrayReference.create(symbol, [int_one.copy(), sum_op2]),
+        one.copy())
+    assert assignment.is_array_range is False
+
+    # lhs is a scalar member of a structure
+    grid_type = StructureType.create([
+        ("dx", REAL_SINGLE_TYPE, Symbol.Visibility.PUBLIC),
+        ("dy", REAL_SINGLE_TYPE, Symbol.Visibility.PUBLIC)])
+    grid_type_symbol = TypeSymbol("grid_type", grid_type)
+    grid_sym = DataSymbol("grid", grid_type_symbol)
+    assignment = Assignment.create(StructureReference.create(grid_sym, ["dx"]),
+                                   one.copy())
+    assert assignment.is_array_range is False