diff --git a/pystencils/transformations.py b/pystencils/transformations.py
index 6411d6acecf78297f551e5f8c479701a281b597b..9924b0a62dfefe311e0f72d5c63558fa6d189690 100644
--- a/pystencils/transformations.py
+++ b/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):
diff --git a/pystencils_tests/test_blocking.py b/pystencils_tests/test_blocking.py
index cc5ca78cbe658b96a69e0168137eb6083719a747..c554da74daefc618850028595c30a9eb1cf2754a 100644
--- a/pystencils_tests/test_blocking.py
+++ b/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])
diff --git a/pystencils_tests/test_move_constant_before_loop.py b/pystencils_tests/test_move_constant_before_loop.py
new file mode 100644
index 0000000000000000000000000000000000000000..a76ab82c2ee914451a005378283598c1bb1991d7
--- /dev/null
+++ b/pystencils_tests/test_move_constant_before_loop.py
@@ -0,0 +1,30 @@
+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