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):
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:
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.
if not isinstance(child, ast.SympyAssignment): # only move SympyAssignments
target, child_to_insert_before = find_block_to_move_to(child)
if target == block: # movement not possible
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)
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
target.insert_before(child, child_to_insert_before)
elif exists_already and exists_already.rhs == child.rhs:
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):
......@@ -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,
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])
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])
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] != loop.parent.args[1]
Supports Markdown
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