def test_printAST(block1, block2, block3, block4):
str1 = printAST(block1)
assert str1 in """
<Iteration i::i::(0, 3, 1)::(0, 0)>
<Iteration j::j::(0, 5, 1)::(0, 0)>
<Iteration k::k::(0, 7, 1)::(0, 0)>
<Expression a[i] = a[i] + b[i] + 5.0>
"""
str2 = printAST(block2)
assert str2 in """
<Iteration i::i::(0, 3, 1)::(0, 0)>
<Expression a[i] = a[i] + b[i] + 5.0>
<Iteration j::j::(0, 5, 1)::(0, 0)>
<Iteration k::k::(0, 7, 1)::(0, 0)>
<Expression a[i] = -a[i] + b[i]>
"""
str3 = printAST(block3)
assert str3 in """
<Iteration i::i::(0, 3, 1)::(0, 0)>
<Iteration s::s::(0, 4, 1)::(0, 0)>
<Expression a[i] = a[i] + b[i] + 5.0>
<Iteration j::j::(0, 5, 1)::(0, 0)>
<Iteration k::k::(0, 7, 1)::(0, 0)>
<Expression a[i] = -a[i] + b[i]>
<Expression a[i] = 4*a[i]*b[i]>
<Iteration q::q::(0, 4, 1)::(0, 0)>
<Expression a[i] = 8.0*a[i] + 6.0/b[i]>
"""
str4 = printAST(block4)
assert str4 in """
def test_nested_transformer(exprs, iters, block2):
"""When created with the kwarg ``nested=True``, a Transformer performs
nested replacements. This test simultaneously replace an inner expression
and an Iteration sorrounding it."""
target_loop = block2.nodes[1]
target_expr = target_loop.nodes[0].nodes[0]
mapper = {target_loop: iters[3](target_loop.nodes[0]),
target_expr: exprs[3]}
processed = Transformer(mapper, nested=True).visit(block2)
assert printAST(processed) == """<Iteration i::i::(0, 3, 1)::(0, 0)>
def test_printAST(block1, block2, block3):
str1 = printAST(block1)
assert str1 in """
<Iteration i::i::(0, 3, 1)::(0, 0)>
<Iteration j::j::(0, 5, 1)::(0, 0)>
<Iteration k::k::(0, 7, 1)::(0, 0)>
<Expression a[i] = a[i] + b[i] + 5.0>
"""
str2 = printAST(block2)
assert str2 in """
<Iteration i::i::(0, 3, 1)::(0, 0)>
<Expression a[i] = a[i] + b[i] + 5.0>
<Iteration j::j::(0, 5, 1)::(0, 0)>
<Iteration k::k::(0, 7, 1)::(0, 0)>
<Expression a[i] = -a[i] + b[i]>
"""
str3 = printAST(block3)
assert str3 in """
def test_nested_transformer(exprs, iters, block2):
"""Unlike Transformer, based on BFS, a NestedTransformer applies transformations
performing a DFS. This test simultaneously replace an inner expression and an
Iteration sorrounding it."""
target_loop = block2.nodes[1]
target_expr = target_loop.nodes[0].nodes[0]
mapper = {
target_loop: iters[3](target_loop.nodes[0]),
target_expr: exprs[3]
}
processed = NestedTransformer(mapper).visit(block2)
assert printAST(processed) == """<Iteration i::i::(0, 3, 1)::(0, 0)>
def test_merge_iterations_flat(exprs, iters):
"""Test outer loop merging on a simple two-level hierarchy:
for i for i
for j \ for j
expr0 === \ expr0
for i === / for k
for k / expr1
expr1
"""
block = [iters[0](iters[1](exprs[0])), iters[0](iters[2](exprs[1]))]
newblock = MergeOuterIterations().visit(block)
newstr = printAST(newblock)
assert newstr == """<Iteration i::i::[0, 3, 1]::[0, 0]>
def test_merge_iterations_deep(exprs, iters):
"""Test outer loop merging on a deep hierarchy:
for i for i
for j for j
expr0 \ expr0
for i === \ for k
for k === / expr0
expr0 / expr1
for k
expr1
"""
block = [
iters[0](iters[1](exprs[0])),
iters[0]([iters[2](exprs[0]), iters[2](exprs[1])])
]
newblock = MergeOuterIterations().visit(block)
newstr = printAST(newblock)
assert newstr == """<Iteration i::i::[0, 3, 1]::[0, 0]>
def test_merge_iterations_nested(exprs, iters):
"""Test outer loop merging on a nested hierarchy that only exposes
the second-level merge after the first level has been performed:
for i for i
for j for j
expr0 \ expr0
for i === \ expr1
for j === / for k
expr1 / expr1
for k
expr1
"""
block = [
iters[0](iters[1](exprs[0])),
iters[0]([iters[1](exprs[1]), iters[2](exprs[1])])
]
newblock = MergeOuterIterations().visit(block)
newstr = printAST(newblock)
assert newstr == """<Iteration i::i::[0, 3, 1]::[0, 0]>
def pprint(node, verbose=True):
"""
Shortcut to pretty print Iteration/Expression trees.
"""
from devito.ir.iet import printAST
print(printAST(node, verbose))
def pprint(node, verbose=True):
"""
Shortcut to pretty print Iteration/Expression trees.
"""
from devito.ir.iet import printAST
print(printAST(node, verbose))
