def find_and_merge_horizontal_loops(root: Node):
copy = root.copy(deep=True)
vertical_loops = eve.FindNodes().by_type(nir.VerticalLoop, copy)
for loop in vertical_loops:
loop = merge_horizontal_loops(loop, _find_merge_candidates(loop))
return copy
def find_and_merge_horizontal_loops(root: Node):
copy = root.copy(deep=True)
vertical_loops = eve.iter_tree(copy).if_isinstance(
nir.VerticalLoop).to_list()
for loop in vertical_loops:
loop = merge_horizontal_loops(loop, _find_merge_candidates(loop))
return copy
def _update_node(node: Node, updated_children: Dict[str, TreeNode]) -> Expr:
# create new node only if children changed
old_children = node.dict(include={*updated_children.keys()})
if any([
old_children[k] != updated_children[k]
for k in updated_children.keys()
]):
return node.copy(update=updated_children)
else:
return node