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