(Towards/Closes? #3463) max parallel assignments

src/psyclone/psyir/transformations/maximal_omp_parallel_region_trans.py

@@ -38,6 +38,7 @@
 from typing import Union
 
 from psyclone.psyir.nodes import (
+        Assignment,
         OMPTaskwaitDirective,
         OMPBarrierDirective,
         OMPSerialDirective,
@@ -47,10 +48,12 @@
         OMPTaskDirective,
         DynamicOMPTaskDirective,
         Node,
+        Reference,
         Schedule
 )
 from psyclone.psyir.transformations.maximal_region_trans import (
-        MaximalRegionTrans)
+        MaximalRegionTrans
+)
 from psyclone.psyir.transformations.omp_parallel_trans import OMPParallelTrans
 from psyclone.utils import transformation_documentation_wrapper
 
@@ -67,6 +70,8 @@ class MaximalOMPParallelRegionTrans(MaximalRegionTrans):
     _transformation = OMPParallelTrans
     _SUB_TRANSFORMATIONS = [OMPParallelTrans]
     # Tuple of statement nodes allowed inside the _transformation
+    # Note that Assignments may also be allowed, but they are a special
+    # case not included in this selection.
     _allowed_contiguous_statements = (
         OMPTaskwaitDirective,
         OMPBarrierDirective,
@@ -86,6 +91,101 @@ class MaximalOMPParallelRegionTrans(MaximalRegionTrans):
         DynamicOMPTaskDirective,
     )
 
+    def _node_allowed(self, node: Node, current_block: list[Node]) -> bool:
+        '''Returns whether the provided node is allowed in the _transformation.
+
+        Nodes defined in the _allowed_contiguous_statements are always allowed
+        to be in the region to be transformed. All other nodes are rejected
+        apart from Assignment nodes. These can be accepted if the lhs is an
+        assignment to a local, scalar variable and there is read access
+        to the symbol outside of the current_block.
+
+        :param node: the candidate node to be in the transformation region.
+        :param current_block: The current block that node would be added into.
+
+        :returns: whether the node is allowed to be in the transformed region.
+        '''
+        if isinstance(node, self._allowed_contiguous_statements):
+            return True
+
+        if not isinstance(node, Assignment):
+            return False
+
+        # If the current_block is empty there's no need to add the Assignment
+        # into the parallel region.
+        if not current_block:
+            return False
+
+        # Assignments are a special case.
+        # If the lhs is an array or a non-local symbol then it is not allowed.
+        if node.lhs.symbol.is_array or not node.lhs.symbol.is_automatic:
+            return False
+
+        # If the lhs symbol appears on the rhs then its read first so we
+        # dissallow it for now. Note that in theory this could be allowed
+        # if its previously been accessed in this region and is private
+        # but we have no ability to determine this in PSyclone at this stage.
+        lhs_sym = node.lhs.symbol
+        for reference in node.rhs.walk(Reference):
+            if reference.symbol.name == lhs_sym.name:
+                return False
+
+        # We know we have an assignment where the lhs is a local scalar
+        # and does not appear on the right hand side. We need to check all
+        # the following read accesses appear in the current_block.
+        accesses = node.lhs.next_accesses()
+        # Find the abs_position of the final node in the current_block.
+        last_element = current_block[-1]
+        # For reviewer - this could be implemented using
+        # following_node.abs_position-1, but if there is no following_node
+        # we'd still need to return to this (or I guess some recursive
+        # loop over children[-1]. Let me know if you have a preference
+        # w.r.t performance as this is probably the worst option but
+        # easy to read.
+        start_position = node.lhs.abs_position
+        end_position = last_element.walk(Node)[-1].abs_position
+
+        for access in accesses:
+            # We can ignore writes as this overwrites any data from
+            # this assignment.
+            if isinstance(access, Reference) and access.is_write:
+                continue
+            pos = access.abs_position
+            # If we have a read access outside of the region then
+            # we can't add this assignment to the region.
+            if pos < start_position or pos > end_position:
+                return False
+        # Privatisation of variables is handled when applying the
+        # transformation.
+        return True
+
+    def _apply_transformation(self, block: list[Node],
+                              **kwargs):
+        '''
+        Applies the transformation defined by self._transformation to the
+        block supplied, with the kwargs provided.
+
+        The default implementation does nothing else, however the
+        functionality is provided so subclasses can perform specific
+        operations (or provide additional specific options) as required for
+        their transformations.
+
+        :param block: The block to apply the transformations to.
+        '''
+        # If we have any assignments directly in the block then we need to
+        # do privatisation of the lhs variable
+        force_private = []
+        if "force_private" in kwargs:
+            force_private.extend(kwargs["force_private"])
+            del kwargs["force_private"]
+        for node in block:
+            if isinstance(node, Assignment):
+                if node.lhs.symbol.name not in force_private:
+                    force_private.append(node.lhs.symbol.name)
+
+        self._transformation().apply(block, force_private=force_private,
+                                     **kwargs)
+
     def apply(self, nodes: Union[Node, Schedule, list[Node]], **kwargs):
         '''Applies the transformation to the nodes provided.
 

src/psyclone/psyir/transformations/maximal_region_trans.py

@@ -76,7 +76,7 @@ class MaximalRegionTrans(RegionTrans, metaclass=abc.ABCMeta):
     #: only barriers doesn't need to be transformed). Defaults to any Node.
     _required_nodes = (Node)
 
-    def _node_allowed(self, node: Node) -> bool:
+    def _node_allowed(self, node: Node, current_block: list[Node]) -> bool:
         '''Returns whether the provided node is allowed in the _transformation.
 
         The default implementation checks whether the node is an instance
@@ -85,6 +85,8 @@ def _node_allowed(self, node: Node) -> bool:
         function is pure).
 
         :param node: the candidate node to be in the transformation region.
+        :param current_block: The current block that node would be added into.
+
         :returns: whether the node is allowed to be in the transformed region.
         '''
         return isinstance(node, self._allowed_contiguous_statements)
@@ -108,24 +110,27 @@ def _satisfies_minimum_region_rules(self, region: list[Node]) -> bool:
                 return True
         return False
 
-    def _can_be_in_region(self, node: Node) -> bool:
+    def _can_be_in_region(
+        self, node: Node, current_block: list[Node]
+    ) -> bool:
         '''Returns whether the provided node can be included in a
         region. Loops and if statements are recursed into to check if their
         children can be.
 
         :param node: the candidate Node to be placed into a transformed
             region.
+        :param current_block: The current block that node would be added into.
 
         :returns: whether it is safe to add the node to a transformed region.
         '''
-        if self._node_allowed(node):
+        if self._node_allowed(node, current_block):
             return True
 
         if isinstance(node, (Loop, WhileLoop)):
             # Check that all contents of the loop body can be part
             # of the region.
             for child in node.loop_body:
-                if not self._can_be_in_region(child):
+                if not self._can_be_in_region(child, current_block):
                     break
             else:
                 return True
@@ -136,11 +141,12 @@ def _can_be_in_region(self, node: Node) -> bool:
             # of the region.
             allowed = True
             for child in node.if_body:
-                allowed = (allowed and self._can_be_in_region(child))
+                allowed = (allowed and self._can_be_in_region(child,
+                                                              current_block))
             if node.else_body and allowed:
                 for child in node.else_body:
                     allowed = (allowed and
-                               self._can_be_in_region(child))
+                               self._can_be_in_region(child, current_block))
             return allowed
 
         # All other node types we default to False.
@@ -164,52 +170,52 @@ def _compute_transformable_sections(
         # Find the largest sections we can surround with the transformation.
         all_blocks = []
         current_block = []
-        for child in node_list:
+        for child in reversed(node_list):
             # If the child can be added to a transformed region then add it
             # to the current block of nodes.
-            if self._can_be_in_region(child):
+            if self._can_be_in_region(child, current_block):
                 # Check that validation still succeeds if we add this child
                 # to the current block.
                 try:
-                    trans.validate(current_block + [child], **trans_kwargs)
-                    current_block.append(child)
+                    trans.validate([child] + current_block, **trans_kwargs)
+                    current_block.insert(0, child)
                 except TransformationError:
                     # If validation now fails, then don't add this to the
                     # current block and add the block to the allowed blocks
                     # if allowed.
                     if current_block:
                         if self._satisfies_minimum_region_rules(current_block):
-                            all_blocks.append(current_block)
+                            all_blocks.insert(0, current_block)
                         current_block = []
             else:
                 # Otherwise, if the current_block contains any children,
                 # add them to the list of regions to be transformed and reset
                 # the current_block.
                 if current_block:
                     if self._satisfies_minimum_region_rules(current_block):
-                        all_blocks.append(current_block)
+                        all_blocks.insert(0, current_block)
                     current_block = []
                 # Need to recurse on some node types
                 if isinstance(child, IfBlock):
-                    if_blocks = self._compute_transformable_sections(
-                            child.if_body, trans, trans_kwargs
-                    )
-                    all_blocks.extend(if_blocks)
                     if child.else_body:
                         else_blocks = self._compute_transformable_sections(
-                            child.else_body, trans, trans_kwargs
+                            child.else_body[:], trans, trans_kwargs
                         )
-                        all_blocks.extend(else_blocks)
+                        all_blocks = else_blocks + all_blocks
+                    if_blocks = self._compute_transformable_sections(
+                            child.if_body[:], trans, trans_kwargs
+                    )
+                    all_blocks = if_blocks + all_blocks
                 if isinstance(child, (Loop, WhileLoop)):
                     loop_blocks = self._compute_transformable_sections(
-                        child.loop_body, trans, trans_kwargs
+                        child.loop_body[:], trans, trans_kwargs
                     )
-                    all_blocks.extend(loop_blocks)
+                    all_blocks = loop_blocks + all_blocks
         # If any nodes are left in the current block at the end of the
         # node_list, then add them to a transformed region
         if current_block:
             if self._satisfies_minimum_region_rules(current_block):
-                all_blocks.append(current_block)
+                all_blocks.insert(0, current_block)
 
         return all_blocks
 
@@ -242,6 +248,21 @@ def validate(self, nodes: Union[Node, Schedule, list[Node]], **kwargs):
                     f"{prev_position}.")
             prev_position = child.position
 
+    def _apply_transformation(self, block: list[Node],
+                              **kwargs):
+        '''
+        Applies the transformation defined by self._transformation to the
+        block supplied, with the kwargs provided.
+
+        The default implementation does nothing else, however the
+        functionality is provided so subclasses can perform specific
+        operations (or provide additional specific options) as required for
+        their transformations.
+
+        :param block: The block to apply the transformations to.
+        '''
+        self._transformation.apply(block, **kwargs)
+
     def apply(self, nodes: Union[Node, Schedule, list[Node]], **kwargs):
         '''Applies the transformation to the nodes provided.
 
@@ -260,4 +281,4 @@ def apply(self, nodes: Union[Node, Schedule, list[Node]], **kwargs):
 
         # Apply the transformation to all of the blocks found.
         for block in all_blocks:
-            par_trans.apply(block, **tr_kwargs)
+            self._apply_transformation(block, **tr_kwargs)

src/psyclone/tests/psyir/transformations/maximal_ompparallel_region_trans_test.py

@@ -37,6 +37,7 @@
 
 from psyclone.psyir.nodes import (
     Loop,
+    Routine,
     OMPBarrierDirective,
     OMPParallelDirective,
 )
@@ -153,3 +154,77 @@ def test_maximal_ompparallel_region_trans_apply(fortran_reader):
     assert isinstance(psyir.children[0].children[0], OMPParallelDirective)
     assert isinstance(psyir.children[0].children[2], OMPBarrierDirective)
     assert isinstance(psyir.children[0].children[3], OMPBarrierDirective)
+
+
+def test_maximal_ompparallel_region_trans_apply_provided_forceprivate(
+    fortran_reader, fortran_writer
+):
+    '''Test that the provided force_private option is correctly passed
+    to the subtransformation.'''
+    psyir = fortran_reader.psyir_from_source('''
+        module my_mod
+            contains
+            subroutine my_subroutine()
+                integer :: ji, jj, jk, jpkm1, jpjm1, jpim1, scalar1
+                real, dimension(10, 10, 10) :: array1, array2
+                array2 = 1
+                do jk = 2, jpkm1, 1
+                  do jj = 2, jpjm1, 1
+                    do ji = 2, jpim1, 1
+                       array2(ji,jj,jk) = array2(ji,jj,jk) + 1
+                       array1(ji,jj,jk) = array2(ji,jj,jk)
+                    enddo
+                  enddo
+                enddo
+            end subroutine
+        end module my_mod''')
+    ltrans = OMPLoopTrans()
+    loops = psyir.walk(Loop)
+    ltrans.apply(loops[0])
+    maxpartrans = MaximalOMPParallelRegionTrans()
+    routine = psyir.walk(Routine)[0]
+
+    maxpartrans.apply(routine, force_private=['array2'])
+    # array2 is requested private, the shared_attibute_inference promotes it to
+    # firstprivate because it is read first
+    expected = '''\
+    !$omp parallel default(shared) private(ji,jj,jk) firstprivate(array2)
+    !$omp do schedule(auto)
+    do jk = 2, jpkm1, 1\n'''
+
+    gen = fortran_writer(psyir)
+    assert expected in gen
+
+
+def test_maximal_ompparallel_region_trans_apply_assignment(
+    fortran_reader, fortran_writer
+):
+    '''Test that assignments to scalars can appear in the
+    parallel region as expected.'''
+    psyir = fortran_reader.psyir_from_source('''
+    subroutine test
+    use omp_lib, only: omp_get_thread_num
+    integer :: tid, i
+    integer, dimension(100) :: arr
+
+    tid = omp_get_thread_num()
+    do i = 1, 100
+        arr(i) = i + tid
+    end do
+    end subroutine test''')
+    ltrans = OMPLoopTrans()
+    loops = psyir.walk(Loop)
+    ltrans.apply(loops[0])
+    maxpartrans = MaximalOMPParallelRegionTrans()
+    routine = psyir.walk(Routine)[0]
+    maxpartrans.apply(routine)
+    out = fortran_writer(psyir)
+    correct = """  !$omp parallel default(shared) private(i,tid)
+  tid = omp_get_thread_num()
+  !$omp do schedule(auto)
+  do i = 1, 100, 1
+    arr(i) = i + tid
+  enddo
+  !$omp end do
+  !$omp end parallel"""
+    assert correct in out