diff --git a/tinygrad/codegen/late/linearizer.py b/tinygrad/codegen/late/linearizer.py index 6c95b37880..a12fd0b744 100644 --- a/tinygrad/codegen/late/linearizer.py +++ b/tinygrad/codegen/late/linearizer.py @@ -64,7 +64,10 @@ class CFGContext: # ranges that have dependencies on other siblings need to be scheduled after them order = sorted(v, key=lambda x: len([u for u in v if u in deps[x]])) zipped = zip(order, order[1:]) if k.op is Ops.SINK else zip([k.src[1]] + order, order) - for x,y in zipped: self.edges[y.src[1]] = x + for x,y in zipped: + # TODO: this can happen! it causes infinite loop in shufflenet + assert y.src[1] not in x.backward_slice_with_self + self.edges[y.src[1]] = x pm_add_control_flow = PatternMatcher([ (UPat(Ops.RANGE, name="x"), lambda ctx,x: x.replace(src=x.src+(y,)) if (y:=ctx.edges.get(x)) is not None else None), diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index a89536f10b..b5dfd4ab95 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -452,7 +452,7 @@ def _valid_priority(v: UOp, valids:list[UOp]): return sum(-1 if (res:=parse_valid(v)) is not None and res[0] in other.toposort() else 0 for other in valids) def simplify_valid(valid:UOp) -> UOp|None: - if valid.op_in_backward_slice_with_self(Ops.LOAD): return None # this should only be for indexing, skip if there's a LOAD + if valid.op_in_backward_slice_with_self(Ops.INDEX): return None # this should only be for indexing, skip if there's a INDEX ret:list[UOp] = [] something_changed = False valids = list(valid.split_uop(Ops.AND))