def _extract_accessors(node: eve.Node) -> List[gtcpp.GTAccessor]:
extents = (
node.iter_tree()
.if_isinstance(gtcpp.AccessorRef)
.reduceby(
(lambda extent, accessor_ref: extent + accessor_ref.offset),
"name",
init=gtcpp.GTExtent.zero(),
as_dict=True,
)
)
inout_fields: Set[str] = (
node.iter_tree()
.if_isinstance(gtcpp.AssignStmt)
.getattr("left")
.if_isinstance(gtcpp.AccessorRef)
.getattr("name")
.to_set()
)
return [
gtcpp.GTAccessor(
name=name,
id=i,
intent=gtcpp.Intent.INOUT if name in inout_fields else gtcpp.Intent.IN,
extent=extent,
)
for i, (name, extent) in enumerate(extents.items())
]
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
def _offset_limit(root: Node) -> int:
return (
root.iter_tree()
.if_isinstance(gtcpp.GTLevel)
.getattr("offset")
.reduce(lambda state, cur: max(state, abs(cur)), init=0)
) + 1
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
