def copy_arrays(mapper, reverse=False):
"""
Build an Iteration/Expression tree performing the copy ``k = v``, or
``v = k`` if reverse=True, for each (k, v) in mapper. (k, v) are expected
to be of type :class:`IndexedData`. The loop bounds are inferred from
the dimensions used in ``k``.
"""
if not mapper:
return ()
# Build the Iteration tree for the copy
iterations = []
for k, v in mapper.items():
handle = []
indices = k.function.indices
for i, j in zip(k.shape, indices):
handle.append(Iteration([], dimension=j, limits=i))
lhs, rhs = (v, k) if reverse else (k, v)
handle.append(
Expression(Eq(lhs[indices], rhs[indices]), dtype=k.function.dtype))
iterations.append(compose_nodes(handle))
# Maybe some Iterations are mergeable
iterations = MergeOuterIterations().visit(iterations)
return iterations
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]>