github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-26 23:38:58 -05:00
Code Issues Packages Projects Releases Wiki Activity

sink folding rule [run_process_replay] (#5145)

Browse Source
This commit is contained in:
George Hotz
2024-06-25 18:34:44 -07:00
committed by GitHub
parent fa78755f19
commit cc7fafcd8b
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tinygrad/codegen/uops.py
View File

@@ -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)