Merge branch 'master' into 2027_effective_shape_bug

changelog

@@ -247,6 +247,8 @@
 	Python used in the CI.
 
 	85) PR #2748. Revert Python used on RTD to 3.12 as 3.13 unsupported.
+
+	86) PR #2707 for #257. Fixes bugs in the loop-fusion transformation.
 
 release 2.5.0 14th of February 2024
 

src/psyclone/core/symbolic_maths.py

@@ -82,13 +82,20 @@ def get():
 
     # -------------------------------------------------------------------------
     @staticmethod
-    def equal(exp1, exp2):
+    def equal(exp1, exp2, identical_variables=None):
         '''Test if the two PSyIR expressions are symbolically equivalent.
+        The optional identical_variables dictionary can contain information
+        about variables which are known to be the same. For example, if
+        `identical_variables={'i': 'j'}`, then 'i+1' and 'j+1' will be
+        considered equal.
 
         :param exp1: the first expression to be compared.
         :type exp1: :py:class:`psyclone.psyir.nodes.Node`
         :param exp2: the second expression to be compared.
         :type exp2: :py:class:`psyclone.psyir.nodes.Node`
+        :param identical_variables: which variable names are known to be
+            identical
+        :type identical_variables: Optional[dict[str, str]]
 
         :returns: whether the two expressions are mathematically \
             identical.
@@ -99,7 +106,8 @@ def equal(exp1, exp2):
         if exp1 is None or exp2 is None:
             return exp1 == exp2
 
-        diff = SymbolicMaths._subtract(exp1, exp2)
+        diff = SymbolicMaths._subtract(exp1, exp2,
+                                       identical_variables=identical_variables)
         # For ranges all values (start, stop, step) must be equal, meaning
         # each index of the difference must evaluate to 0:
         if isinstance(diff, list):
@@ -164,17 +172,25 @@ def never_equal(exp1, exp2):
 
     # -------------------------------------------------------------------------
     @staticmethod
-    def _subtract(exp1, exp2):
+    def _subtract(exp1, exp2, identical_variables=None):
         '''Subtracts two PSyIR expressions and returns the simplified result
         of this operation. An expression might result in multiple SymPy
         expressions - for example, a `Range` node becomes a 3-tuple (start,
         stop, step). In this case, each of the components will be handled
         individually, and a list will be returned.
 
+        The optional identical_variables dictionary can contain information
+        about variables which are known to be the same. For example, if
+        `identical_variables={'i': 'j'}`, then 'i+1' and 'j+1' will be
+        considered equal.
+
         :param exp1: the first expression to be compared.
         :type exp1: Optional[:py:class:`psyclone.psyir.nodes.Node`]
         :param exp2: the second expression to be compared.
         :type exp2: Optional[:py:class:`psyclone.psyir.nodes.Node`]
+        :param identical_variables: which variable names are known to be
+            identical
+        :type identical_variables: Optional[dict[str, str]]
 
         :returns: the sympy expression resulting from subtracting exp2 \
             from exp1.
@@ -188,7 +204,10 @@ def _subtract(exp1, exp2):
 
         # Use the SymPyWriter to convert the two expressions to
         # SymPy expressions:
-        sympy_expressions = SymPyWriter(exp1, exp2)
+        writer = SymPyWriter()
+        sympy_expressions = writer([exp1, exp2],
+                                   identical_variables=identical_variables)
+
         # If an expression is a range node, then the corresponding SymPy
         # expression will be a tuple:
         if isinstance(sympy_expressions[0], tuple) and \

src/psyclone/psyir/backend/sympy_writer.py

@@ -246,7 +246,7 @@ def _create_sympy_array_function(self, name, sig=None, num_dims=None):
         return new_func
 
     # -------------------------------------------------------------------------
-    def _create_type_map(self, list_of_expressions):
+    def _create_type_map(self, list_of_expressions, identical_variables=None):
         '''This function creates a dictionary mapping each Reference in any
         of the expressions to either a SymPy Function (if the reference
         is an array reference) or a Symbol (if the reference is not an
@@ -269,6 +269,15 @@ def _create_type_map(self, list_of_expressions):
         ``lambda_1``). The SymPyWriter (e.g. in ``reference_node``) will
         use the renamed value when creating the string representation.
 
+        The optional identical_variables dictionary can contain information
+        about variables which are known to be the same. For example, if
+        `identical_variables={'i': 'j'}`, then 'i+1' and 'j+1' will be
+        considered equal.
+
+        :param identical_variables: which variable names are known to be
+            identical
+        :type identical_variables: Optional[dict[str, str]]
+
         :param list_of_expressions: the list of expressions from which all
             references are taken and added to a symbol table to avoid
             renaming any symbols (so that only member names will be renamed).
@@ -320,6 +329,16 @@ def _create_type_map(self, list_of_expressions):
                 self._sympy_type_map[unique_sym.name] = \
                     self._create_sympy_array_function(name)
 
+        if not identical_variables:
+            identical_variables = {}
+        # For all variables that are the same, set the symbols to be
+        # identical. This means if e.g. identical_variables={'i': 'j'},
+        # the expression i-j becomes j-j = 0
+
+        for var1, var2 in identical_variables.items():
+            if var1 in self._sympy_type_map and var2 in self._sympy_type_map:
+                self._sympy_type_map[var1] = self._sympy_type_map[var2]
+
         # Now all symbols have been added to the symbol table, create
         # unique names for the lower- and upper-bounds using special tags:
         self._lower_bound = \
@@ -358,11 +377,19 @@ def type_map(self):
         return self._sympy_type_map
 
     # -------------------------------------------------------------------------
-    def _to_str(self, list_of_expressions):
+    def _to_str(self, list_of_expressions, identical_variables=None):
         '''Converts PSyIR expressions to strings. It will replace Fortran-
         specific expressions with code that can be parsed by SymPy. The
         argument can either be a single element (in which case a single string
         is returned) or a list/tuple, in which case a list is returned.
+        The optional identical_variables dictionary can contain information
+        about variables which are known to be the same. For example, if
+        `identical_variables={'i': 'j'}`, then 'i+1' and 'j+1' will be
+        considered equal.
+
+        :param identical_variables: which variable names are known to be
+            identical
+        :type identical_variables: Optional[dict[str, str]]
 
         :param list_of_expressions: the list of expressions which are to be
             converted into SymPy-parsable strings.
@@ -379,7 +406,8 @@ def _to_str(self, list_of_expressions):
 
         # Create the type map in `self._sympy_type_map`, which is required
         # when converting these strings to SymPy expressions
-        self._create_type_map(list_of_expressions)
+        self._create_type_map(list_of_expressions,
+                              identical_variables=identical_variables)
 
         expression_str_list = []
         for expr in list_of_expressions:
@@ -392,19 +420,26 @@ def _to_str(self, list_of_expressions):
         return expression_str_list
 
     # -------------------------------------------------------------------------
-    def __call__(self, list_of_expressions):
+    def __call__(self, list_of_expressions, identical_variables=None):
         '''
         This function takes a list of PSyIR expressions, and converts
         them all into Sympy expressions using the SymPy parser.
         It takes care of all Fortran specific conversion required (e.g.
         constants with kind specification, ...), including the renaming of
         member accesses, as described in
         https://psyclone-dev.readthedocs.io/en/latest/sympy.html#sympy
+        The optional identical_variables dictionary can contain information
+        about variables which are known to be the same. For example, if
+        `identical_variables={'i': 'j'}`, then 'i+1' and 'j+1' will be
+        considered equal.
 
         :param list_of_expressions: the list of expressions which are to be
             converted into SymPy-parsable strings.
         :type list_of_expressions: list of
             :py:class:`psyclone.psyir.nodes.Node`
+        :param identical_variables: which variable names are known to be
+            identical
+        :type identical_variables: Optional[dict[str, str]]
 
         :returns: a 2-tuple consisting of the the converted PSyIR
             expressions, followed by a dictionary mapping the symbol names
@@ -413,12 +448,25 @@ def __call__(self, list_of_expressions):
                       List[:py:class:`sympy.core.basic.Basic`]]
 
         :raises VisitorError: if an invalid SymPy expression is found.
+        :raises TypeError: if the identical_variables parameter is not
+            a dict, or does contain a key or value that is not a string.
 
         '''
+        if identical_variables:
+            if not isinstance(identical_variables, dict):
+                raise TypeError(f"Expected identical_variables to be "
+                                f"a dictionary, but got "
+                                f"{type(identical_variables)}.")
+            if any(not isinstance(key, str) or not isinstance(value, str)
+                   for key, value in identical_variables.items()):
+                raise TypeError("Dictionary identical_variables "
+                                "contains a non-string key or value.")
+
         is_list = isinstance(list_of_expressions, (tuple, list))
         if not is_list:
             list_of_expressions = [list_of_expressions]
-        expression_str_list = self._to_str(list_of_expressions)
+        expression_str_list = self._to_str(
+            list_of_expressions, identical_variables=identical_variables)
 
         result = []
         for expr in expression_str_list: