diff --git a/tinygrad/codegen/uops.py b/tinygrad/codegen/uops.py index 2078ed2d4d..d7bbd5003c 100644 --- a/tinygrad/codegen/uops.py +++ b/tinygrad/codegen/uops.py @@ -271,6 +271,9 @@ constant_folder = PatternMatcher([ (UOp.load(UOp.var(), UOp.var(), UOp.const(None, 0), UOp.cvar("var"), UOp.var()), lambda var: var), (UOp.store(UOp.var("buf"), UOp.var("idx"), UOp.var("val"), UOp.const(None, 1)), UOp.store), (UOp.store(UOp.var(), UOp.var(), UOp.var(), UOp.const(None, 0)), lambda: UOp(UOps.NOOP)), + # remove NOOPs from SINK + (UPat(UOps.SINK, name="root"), + lambda root: UOp(UOps.SINK, root.dtype, a, root.arg) if len(a:=tuple(x for x in root.src if x.op is not UOps.NOOP)) != len(root.src) else None) ]) # *** uop graph *** @@ -383,7 +386,6 @@ class UOpGraph: graph[x].append(u) if u.op is UOps.RANGE: loops.append(u) if u.op is UOps.IF: ifs.append(u) - sink = UOp(UOps.SINK, None, tuple(x for x in sink.src if x.op is not UOps.NOOP)) add_parents(sink) @functools.lru_cache(None)