diff --git a/tinygrad/codegen/uops.py b/tinygrad/codegen/uops.py index 1c4010c332..f388c0c707 100644 --- a/tinygrad/codegen/uops.py +++ b/tinygrad/codegen/uops.py @@ -289,7 +289,7 @@ class UOpGraph: for i,u in enumerate(self): print(f"{i:4d} {str(u.op):20s}: {str(u.dtype) if u.dtype is not None else '':25s} " f"{str([self.uops.index(x) for x in u.src]):32s} {u.arg}") - def graph_rewrite(self, sink, pm): + def graph_rewrite(self, sink:UOp, pm:PatternMatcher): # recursive rewrite changed = getenv("UOPS_REWRITE", 1) run_cnt = 0 @@ -317,7 +317,7 @@ class UOpGraph: assert run_cnt < 100, "exceeded 100 rewrite loops!" return sink - def graph_dedup(self, sink): + def graph_dedup(self, sink:UOp): # add nodes to graph in reverse BFS order # dedup all nodes # TODO: i feel like this BFS is written in a few places, possible to library it? @@ -363,8 +363,8 @@ class UOpGraph: # BFS toposort graph: DefaultDict[UOp, List[UOp]] = defaultdict(list) in_degree: DefaultDict[UOp, int] = defaultdict(int) - loops = [] - ifs = [] + loops:List[UOp] = [] + ifs:List[UOp] = [] nodes: Dict[UOp, None] = {} def add_parents(u:UOp): if u in nodes: return @@ -386,8 +386,8 @@ class UOpGraph: end_for_uop = {UOps.IF:(UOps.STORE, UOps.ENDIF), UOps.RANGE:(UOps.PHI, UOps.ENDRANGE)} scope_children = {p:get_recursive_children(p, end_for_uop[p.op][0]) for p in (loops+ifs)[::-1]} - queue: List = [] - def push(u): + queue:List[Tuple[int, UOp]] = [] + def push(u:UOp): priority = 0 # prefer uops that are loop children for l, ss in scope_children.items():