github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-02-10 06:35:33 -05:00
Code Issues Packages Projects Releases Wiki Activity

Restore fast path for matching new_src in rewrite (#6870)

Browse Source
This commit is contained in:
Tim Becker
2024-10-03 23:22:24 -04:00
committed by GitHub
parent 8931f20765
commit d42cb5596f
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tinygrad/ops.py
View File

@@ -607,8 +607,8 @@ class RewriteContext:
def rewrite(self, n:UOp) -> UOp:
if (rn := self.replace.get(n)) is not None: return rn
new_src = tuple(map(self.rewrite, n.src))
x = UOp(n.op, n.dtype, new_src, n.arg) if new_src != n.src else n
self.replace[n] = ret = self.rewrite(new_x) if (new_x := self.pm.rewrite(x, self.ctx)) is not None else x
new_n = self.pm.rewrite(n, self.ctx) if new_src == n.src else UOp(n.op, n.dtype, new_src, n.arg)
self.replace[n] = ret = n if new_n is None else self.rewrite(new_n)
return ret
def graph_rewrite(sink:UOp, pm:PatternMatcher, ctx=None) -> UOp: