Commit eec4dc4b authored by Martin Bauer's avatar Martin Bauer
Browse files

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
parent f504b40f
...@@ -560,11 +560,11 @@ def move_constants_before_loop(ast_node): ...@@ -560,11 +560,11 @@ def move_constants_before_loop(ast_node):
element = element.parent element = element.parent
return last_block, last_block_child 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: for arg in target_block.args:
if type(arg) is not ast.SympyAssignment: if type(arg) is not ast.SympyAssignment:
continue 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 arg
return None return None
...@@ -579,22 +579,45 @@ def move_constants_before_loop(ast_node): ...@@ -579,22 +579,45 @@ def move_constants_before_loop(ast_node):
get_blocks(ast_node, all_blocks) get_blocks(ast_node, all_blocks)
for block in all_blocks: for block in all_blocks:
children = block.take_child_nodes() 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: 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) target, child_to_insert_before = find_block_to_move_to(child)
if target == block: # movement not possible if target == block: # movement not possible
target.append(child) target.append(child)
else: else:
if isinstance(child, ast.SympyAssignment): 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: else:
exists_already = False exists_already = False
if not exists_already: 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: elif exists_already and exists_already.rhs == child.rhs:
pass pass
else: 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): def split_inner_loop(ast_node: ast.Node, symbol_groups):
......
...@@ -14,7 +14,7 @@ def jacobi(dst, src): ...@@ -14,7 +14,7 @@ def jacobi(dst, src):
def check_equivalence(assignments, src_arr): def check_equivalence(assignments, src_arr):
for openmp in (False, True): 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, with_blocking = ps.create_kernel(assignments, cpu_blocking=(8, 16, 4), cpu_openmp=openmp,
cpu_vectorize_info=vectorization).compile() cpu_vectorize_info=vectorization).compile()
without_blocking = ps.create_kernel(assignments).compile() without_blocking = ps.create_kernel(assignments).compile()
...@@ -28,7 +28,7 @@ def check_equivalence(assignments, src_arr): ...@@ -28,7 +28,7 @@ def check_equivalence(assignments, src_arr):
def test_jacobi3d_var_size(): 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") print("Var Size: Smaller than block sizes")
arr = np.empty([4, 5, 6]) arr = np.empty([4, 5, 6])
......
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
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment