# -----------------------------------------------------------------------------
# BSD 3-Clause License
# Copyright (c) 2020-2026, Science and Technology Facilities Council.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# * Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
# -----------------------------------------------------------------------------
# Author R. W. Ford, N. Nobre and S. Siso, STFC Daresbury Lab
# Modified by J. Henrichs, Bureau of Meteorology
# Modified by A. B. G. Chalk, STFC Daresbury Lab
'''Module providing a transformation from a PSyIR Array Assignment to explicit
PSyIR Loops. This could be useful for e.g. performance reasons, to enable
further transformations e.g. loop fusion or if the back-end does
not support array ranges.
'''
from typing import Optional, Any
from psyclone.errors import LazyString
from psyclone.psyGen import Transformation
from psyclone.psyir.nodes import (
Assignment, Call, IntrinsicCall, Loop, Literal, Node, Range, Reference,
CodeBlock, Routine, BinaryOperation)
from psyclone.psyir.nodes.array_mixin import ArrayMixin
from psyclone.psyir.nodes.structure_accessor_mixin import (
StructureAccessorMixin)
from psyclone.psyir.symbols import (
DataSymbol, ScalarType, SymbolError)
from psyclone.utils import transformation_documentation_wrapper
from psyclone.psyir.transformations.transformation_error import (
TransformationError)
from psyclone.psyir.transformations.reference2arrayrange_trans import (
Reference2ArrayRangeTrans)
[docs]
@transformation_documentation_wrapper
class ArrayAssignment2LoopsTrans(Transformation):
'''Provides a transformation from a PSyIR Array Range to a PSyIR
Loop. For example:
>>> from psyclone.psyir.frontend.fortran import FortranReader
>>> from psyclone.psyir.nodes import Assignment
>>> from psyclone.psyir.transformations import ArrayAssignment2LoopsTrans
>>> code = """
... subroutine sub()
... real :: tmp(10)
... tmp(:) = tmp(:) + 3
... end subroutine sub"""
>>> psyir = FortranReader().psyir_from_source(code)
>>> assignment = psyir.walk(Assignment)[0]
>>> trans = ArrayAssignment2LoopsTrans()
>>> trans.apply(assignment)
>>> print(psyir.debug_string())
subroutine sub()
real, dimension(10) :: tmp
integer :: idx
<BLANKLINE>
do idx = LBOUND(tmp, dim=1), UBOUND(tmp, dim=1), 1
tmp(idx) = tmp(idx) + 3
enddo
<BLANKLINE>
end subroutine sub
<BLANKLINE>
By default the transformation will reject character arrays, though this
can be overridden by setting the 'allow_strings' option to True. Note that
PSyclone expresses syntax such as `character(LEN=100)` as
UnsupportedFortranType, and this transformation will convert unknown or
unsupported types to loops.
'''
[docs]
def apply(
self,
node: Assignment,
options: Optional[dict[str, Any]] = None,
allow_strings: bool = False,
verbose: bool = False,
**kwargs
):
'''Apply the transformation to the specified array assignment node.
Each range node within the assignment is replaced with an explicit
loop. The bounds of the loop are determined from the bounds of the
array range on the left hand side of the assignment.
:param node: an Assignment node.
:param allow_strings: whether to allow the
transformation on a character type array range. Defaults to False.
:param verbose: log the reason the validation failed,
at the moment with a comment in the provided PSyIR node.
:param options: a dictionary with options for transformations.
'''
self.validate(node, options, allow_strings=allow_strings,
verbose=verbose, **kwargs)
# If possible use the routine-level symbol table for nicer names
if node.ancestor(Routine):
symbol_table = node.ancestor(Routine).symbol_table
else:
# We checked in the validate that at least there is a scope
symbol_table = node.scope.symbol_table
# If there is any array reference without the accessor syntax,
# we need to add it first.
for reference in node.walk(Reference):
Reference2ArrayRangeTrans().apply(reference)
# Now replace range expressions with indices for every range in a
# top-level expression (e.g. we don't walk inside non-elemental calls
# because they need to receive the same complete slices). Start with
# the rightmost range to respect Fortran iteration order
for lhs_range in reversed(
_walk_until_non_elemental_call(node.lhs, Range)
):
lhs_array = lhs_range.parent
lhs_index = lhs_array.index_of(lhs_range)
# Create a new, unique, iteration variable for the new loop
loop_variable_symbol = symbol_table.new_symbol(
root_name="idx",
symbol_type=DataSymbol,
datatype=ScalarType.integer_type())
# Replace one range for each top-level array expression in the
# assignment
for expr in reversed(
_walk_until_non_elemental_call(node, ArrayMixin, ArrayMixin)
):
# We use a "for" to capture the value of the last range or
# continue if there is none, but we don't iterate all Ranges
for range_to_replace in reversed(
_walk_until_non_elemental_call(expr, Range)
):
array = range_to_replace.parent
index = array.index_of(range_to_replace)
if lhs_array.same_range(lhs_index, array, index):
# If they iterate over the same bounds we just need a
# reference to the iteration index
index_expr = Reference(loop_variable_symbol)
else:
# If we can not prove that both ranges iterate over the
# exact same bounds we need to provide an offset like:
# array2(idx1 + (lbound_array2 - lbound_array1)
offset = BinaryOperation.create(
BinaryOperation.Operator.SUB,
range_to_replace.start.copy(),
lhs_range.start.copy())
index_expr = BinaryOperation.create(
BinaryOperation.Operator.ADD,
Reference(loop_variable_symbol),
offset)
range_to_replace.replace_with(index_expr)
break # We just substitute one per top-level array
# Replace the assignment with the new explicit loop structure
start, stop, step = lhs_range.pop_all_children()
loop = Loop.create(loop_variable_symbol, start, stop, step, [])
node.replace_with(loop)
loop.loop_body.addchild(node)
def __str__(self):
return ("Convert a PSyIR assignment with array Ranges into explicit "
"PSyIR Loops.")
[docs]
def validate(
self,
node: Assignment,
options: Optional[dict[str, Any]] = None,
**kwargs
):
'''Perform various checks to ensure that it is valid to apply the
ArrayAssignment2LoopsTrans transformation to the supplied PSyIR Node.
By default the validate function will throw an TransformationError
on character arrays, though this can be overridden by setting the
allow_strings option to True. Note that PSyclone expresses syntax such
as `character(LEN=100)` as UnsupportedFortranType, and this
transformation will convert unknown or unsupported types to loops.
:param node: the assignment to transform.
:param options: a dictionary with options for transformations.
:raises TransformationError: if the node argument is not an
Assignment.
:raises TransformationError: if the node argument is an
Assignment whose left hand side is not an ArrayReference.
:raises TransformationError: if the node argument is an
Assignment whose left hand side is an ArrayReference that does
not have Range specifying the access to at least one of its
dimensions.
:raises TransformationError: if two or more of the loop ranges
in the assignment are different or are not known to be the same.
:raises TransformationError: if node contains a character type
child and the allow_strings option is not set.
'''
super().validate(node, **kwargs)
if not options:
self.validate_options(**kwargs)
allow_strings = self.get_option("allow_strings", **kwargs)
verbose = self.get_option("verbose", **kwargs)
else:
# TODO #2668: Deprecate options dictionary.
verbose = options.get("verbose", False)
allow_strings = options.get("allow_strings", False)
if not isinstance(node, Assignment):
raise TransformationError(
f"Error in {self.name} transformation. The supplied node "
f"should be a PSyIR Assignment, but found "
f"'{type(node).__name__}'.")
# Do not allow to transform expressions with CodeBlocks
if node.has_descendant(CodeBlock):
message = (f"{self.name} does not support array assignments that"
f" contain a CodeBlock anywhere in the expression")
if verbose:
node.append_preceding_comment(message)
raise TransformationError(LazyString(
lambda: f"{message}, but found:\n{node.debug_string()}"))
try:
node.scope
except SymbolError:
# pylint: disable=raise-missing-from
raise TransformationError(LazyString(
lambda: f"Error in {self.name} transformation. The "
f"assignment should be in a scope to create the necessary new "
f"symbols, but '{node.debug_string()}' is not."))
node_copy = node.copy()
for reference in node_copy.walk(Reference):
if isinstance(reference.parent, Call):
# These are ok, what is important is what the call returns
# which is checked later
continue
try:
Reference2ArrayRangeTrans().apply(reference)
except TransformationError as err:
message = (
f"Reference2ArrayRangeTrans could not decide if the "
f"reference '{reference.name}' is an array or not. "
f"Resolving the import that brings this variable into "
f"scope may help.")
if verbose:
node.append_preceding_comment(message)
raise TransformationError(message) from err
# After a successful Reference2ArrayRangeTrans all arrays are
# guaranteed to be expressed as ArrayMixin's
array_accessors = node_copy.lhs.walk(ArrayMixin)
if not (isinstance(node_copy.lhs, Reference) and array_accessors):
raise TransformationError(LazyString(
lambda: f"Error in {self.name} transformation. The LHS of the"
f" supplied Assignment node should be a Reference that "
f"contains an array accessor somewhere in the expression, "
f"but found '{node.lhs.debug_string()}'."))
# There should be at least one Range in the LHS
if not node_copy.lhs.walk(Range):
raise TransformationError(LazyString(
lambda: f"Error in {self.name} transformation. The LHS of"
f" the supplied Assignment node should contain an array "
f"accessor with at least one of its dimensions being a "
f"Range, but none were found in '{node.debug_string()}'."))
# All the ArrayMixins must have the same number of Ranges to expand
num_of_ranges = None
for accessor in node_copy.walk(ArrayMixin):
count = len([x for x in accessor.indices if isinstance(x, Range)])
if count > 0:
if not num_of_ranges:
# If it's the first value we find, we store it
num_of_ranges = count
else:
# Otherwise we compare it against the previous found value
if count != num_of_ranges:
raise TransformationError(LazyString(
lambda: f"{self.name} does not support statements"
f" containing array accesses that have varying "
f"numbers of ranges in their accessors, but found:"
f"\n{node.debug_string()}"))
# Check if there is any dependency between the written reference and
# any other reference in the assignment. The check is only against one
# write (assignment lhs top reference) because we fail validation if
# there is any impure call (which could also contain other writes)
written_ref = node_copy.lhs
written_sig, written_idxs = written_ref.get_signature_and_indices()
for ref in node_copy.walk(Reference)[1:]:
if ref.symbol is written_ref.symbol:
found_dependency = False
# If the reference data is not read (e.g. used by an inquiry
# intrinsic), then there is not dependency
if not ref.is_read:
continue
# The accessor is different, there is no dependency
ref_sig, ref_idxs = ref.get_signature_and_indices()
if ref_sig != written_sig:
continue
# If its inside a non-elemental call, we assume it reads the
# whole array, and therefore causes a dependency
parent_call = ref.ancestor(Call)
if parent_call and not parent_call.is_elemental:
found_dependency = True
# The signature indices are provided in a doubly nested
# tuple with components and indices, we need both loops
# below to compare each matching index in both references
for c1, c2 in zip(ref_idxs, written_idxs):
for i1, i2 in zip(c1, c2):
if isinstance(i1, Range) or isinstance(i2, Range):
# If an index is a range, check that the bounds
# are exactly the same or it could be a
# dependency
if i1 != i2:
found_dependency = True
break
# If none of the matching indices are ranges, this
# will be an invariant and not represent a problem
if found_dependency:
raise TransformationError(LazyString(
lambda: f"{self.name} does not support statements "
f"containing dependencies that would generate "
f"loop-carried dependencies when naively "
f"converting them to a loop, but found:"
f"\n{node.debug_string()}"))
# We don't support nested range expressions anywhere in the assignment
# (but note that non-elemental calls reset the expression)
for range_expr in node_copy.walk(Range, stop_type=Range):
# Test that there are no Ranges in any children
# e.g: array(:<cannot have ranges>)
test_nodes = range_expr.children[:]
# or in the member expression if it is a derived type accessor
# e.g: array(:)%<cannot have ranges>
if isinstance(range_expr.parent, StructureAccessorMixin):
test_nodes.append(range_expr.parent.member)
for test in test_nodes:
if _walk_until_non_elemental_call(test, Range):
message = (
f"{self.name} does not support array assignments "
f"that contain nested Range expressions")
if verbose:
node.append_preceding_comment(message)
# pylint: disable=cell-var-from-loop
raise TransformationError(LazyString(lambda: (
f"{message}, but found:\n"
f"{range_expr.parent.debug_string()}")))
# If we allow string arrays then we can skip the check.
if not allow_strings:
message = (f"{self.name} does not expand ranges "
f"on character arrays by default (use the"
f"'allow_strings' option to expand them)")
self.validate_no_char(node, message, verbose)
# We don't accept calls that are not guaranteed to return a scalar
# or be elemental
for call in node_copy.rhs.walk(Call):
if call.is_elemental:
continue
if call.is_pure and isinstance(call.datatype, ScalarType):
continue
if isinstance(call, IntrinsicCall):
name = call.intrinsic.name
else:
name = call.routine.symbol.name
message = (f"{self.name} does not accept calls which are not"
f" guaranteed to return a scalar or be elemental, "
f"but found '{name}'")
if verbose:
node.append_preceding_comment(message)
# pylint: disable=cell-var-from-loop
raise TransformationError(LazyString(
lambda: f"{message} in '{node.debug_string().strip()}'"))
[docs]
@staticmethod
def validate_no_char(node: Node, message: str, verbose: bool) -> None:
'''
Check that there is no character variable accessed in the sub-tree with
the supplied node at its root.
:param node: the root node to check for character assignments.
:param message: the message to use if a character assignment is found.
:param verbose: whether log the reason the validation failed,
at the moment with a comment in the provided PSyIR node.
:raises TransformationError: if the supplied node contains a
child of character type.
'''
for child in node.walk((Literal, Reference)):
if child.is_character(unknown_as=False):
if verbose:
node.append_preceding_comment(message)
# pylint: disable=cell-var-from-loop
raise TransformationError(LazyString(
lambda: f"{message}, but found:"
f"\n{node.debug_string()}"))
def _walk_until_non_elemental_call(
node: Node, search_type: type, stop_type: Optional[type] = None
) -> list[Node]:
''' Custom walk operation that stops at Calls that are not elemental. The
elemental restriction was not possible to implement with the generic
node.walk(), which only supports stopping at given types.
:param node: walk start location.
:param search_type: the class for which the instances are collected.
:param stop_type: the class in which the walk stops (in addition to
non-elemental calls).
:returns: list with all nodes that are instances of search_type
starting at and including the given node.
'''
local_list = []
if isinstance(node, search_type):
local_list.append(node)
if stop_type and isinstance(node, stop_type):
return local_list
if isinstance(node, Call) and not node.is_elemental:
return local_list
for child in node.children:
local_list += _walk_until_non_elemental_call(
child, search_type, stop_type)
return local_list
__all__ = [
'ArrayAssignment2LoopsTrans']
