From eec4dc4b968f91f4e64e0227eafa3c1e9387c243 Mon Sep 17 00:00:00 2001
From: Martin Bauer <martin.bauer@fau.de>
Date: Fri, 26 Apr 2019 14:31:33 +0200
Subject: [PATCH] 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 | 33 ++++++++++++++++---
pystencils_tests/test_blocking.py | 4 +--
.../test_move_constant_before_loop.py | 30 +++++++++++++++++
3 files changed, 60 insertions(+), 7 deletions(-)
create mode 100644 pystencils_tests/test_move_constant_before_loop.py
diff --git a/pystencils/transformations.py b/pystencils/transformations.py
index 6411d6a..9924b0a 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 cc5ca78..c554da7 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 0000000..a76ab82
--- /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
--
GitLab