Enhancement in move_constants_before loop

When two loops have assignments to the same symbol with different
rhs and both are pulled before the loops, one of them is now renamed.
Previously one of them was left inside the loop.

Fixes #27

pystencils/transformations.py

@@ -560,11 +560,11 @@ def move_constants_before_loop(ast_node):
             element = element.parent
         return last_block, last_block_child
 
-    def check_if_assignment_already_in_block(assignment, target_block):
+    def check_if_assignment_already_in_block(assignment, target_block, rhs_or_lhs=True):
         for arg in target_block.args:
             if type(arg) is not ast.SympyAssignment:
                 continue
-            if arg.lhs == assignment.lhs:
+            if (rhs_or_lhs and arg.rhs == assignment.rhs) or (not rhs_or_lhs and arg.lhs == assignment.lhs):
                 return arg
         return None
 
@@ -579,22 +579,45 @@ def move_constants_before_loop(ast_node):
     get_blocks(ast_node, all_blocks)
     for block in all_blocks:
         children = block.take_child_nodes()
+        # Every time a symbol can be replaced in the current block because the assignment
+        # was found in a parent block, but with a different lhs symbol (same rhs)
+        # the outer symbol is inserted here as key.
+        substitute_variables = {}
         for child in children:
+            # Before traversing the next child, all symbols are substituted first.
+            child.subs(substitute_variables)
+
+            if not isinstance(child, ast.SympyAssignment):  # only move SympyAssignments
+                block.append(child)
+                continue
+
             target, child_to_insert_before = find_block_to_move_to(child)
             if target == block:     # movement not possible
                 target.append(child)
             else:
                 if isinstance(child, ast.SympyAssignment):
-                    exists_already = check_if_assignment_already_in_block(child, target)
+                    exists_already = check_if_assignment_already_in_block(child, target, False)
                 else:
                     exists_already = False
 
                 if not exists_already:
-                    target.insert_before(child, child_to_insert_before)
+                    rhs_identical = check_if_assignment_already_in_block(child, target, True)
+                    if rhs_identical:
+                        # there is already an assignment out there with the same rhs
+                        # -> replace all lhs symbols in this block with the lhs of the outer assignment
+                        # -> remove the local assignment (do not re-append child to the former block)
+                        substitute_variables[child.lhs] = rhs_identical.lhs
+                    else:
+                        target.insert_before(child, child_to_insert_before)
                 elif exists_already and exists_already.rhs == child.rhs:
                     pass
                 else:
-                    block.append(child)  # don't move in this case - better would be to rename symbol
+                    # this variable already exists in outer block, but with different rhs
+                    # -> symbol has to be renamed
+                    assert isinstance(child.lhs, TypedSymbol)
+                    new_symbol = TypedSymbol(sp.Dummy().name, child.lhs.dtype)
+                    target.insert_before(ast.SympyAssignment(new_symbol, child.rhs), child_to_insert_before)
+                    substitute_variables[child.lhs] = new_symbol
 
 
 def split_inner_loop(ast_node: ast.Node, symbol_groups):

pystencils_tests/test_blocking.py

@@ -14,7 +14,7 @@ def jacobi(dst, src):
 
 def check_equivalence(assignments, src_arr):
     for openmp in (False, True):
-        for vectorization in (True, False):
+        for vectorization in [False, {'assume_inner_stride_one': True}]:
             with_blocking = ps.create_kernel(assignments, cpu_blocking=(8, 16, 4), cpu_openmp=openmp,
                                              cpu_vectorize_info=vectorization).compile()
             without_blocking = ps.create_kernel(assignments).compile()
@@ -28,7 +28,7 @@ def check_equivalence(assignments, src_arr):
 
 
 def test_jacobi3d_var_size():
-    src, dst = ps.fields("src, dst: double[3D]")
+    src, dst = ps.fields("src, dst: double[3D]", layout='c')
 
     print("Var Size: Smaller than block sizes")
     arr = np.empty([4, 5, 6])

pystencils_tests/test_move_constant_before_loop.py

+import pystencils as ps
+import numpy as np
+from pystencils.astnodes import LoopOverCoordinate, Block, SympyAssignment, TypedSymbol
+from pystencils.transformations import move_constants_before_loop
+
+
+def test_symbol_renaming():
+    """When two loops have assignments to the same symbol with different rhs and both
+    are pulled before the loops, one of them has to be renamed
+    """
+
+    f, g = ps.fields("f, g : double[2D]")
+    a, b, c = [TypedSymbol(n, np.float64) for n in ('a', 'b', 'c')]
+
+    loop1 = LoopOverCoordinate(Block([SympyAssignment(c, a + b),
+                                      SympyAssignment(g[0, 0], f[0, 0] + c)]),
+                               0, 0, 10)
+    loop2 = LoopOverCoordinate(Block([SympyAssignment(c, a ** 2 + b ** 2),
+                                      SympyAssignment(g[0, 0], f[0, 0] + c)]),
+                               0, 0, 10)
+    block = Block([loop1, loop2])
+
+    move_constants_before_loop(block)
+
+    loops = block.atoms(LoopOverCoordinate)
+    assert len(loops) == 2
+    for loop in loops:
+        assert len(loop.body.args) == 1
+        assert len(loop.parent.args) == 4  # 2 loops + 2 subexpressions
+        assert loop.parent.args[0].lhs.name != loop.parent.args[1].lhs.name