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

test

Browse Source
This commit is contained in:
qazal
2024-07-22 14:15:03 +03:00
parent dc21e63bd2
commit 3dc2b48042
1 changed files with 2 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
tinygrad/codegen/uopgraph.py
View File

@@ -478,9 +478,8 @@ class UOpGraph:
sink_srcs = list(self.sink.src)
for i, s in enumerate(sink_srcs):
# breaks for WMMA
if all(x.op is not UOps.WMMA for x in s.parents):
if s.op is UOps.STORE and len(s.src) == 4 and (rw:=_dfs(s, s.src[3])) != s:
sink_srcs[i] = UOp(rw.op, rw.dtype, rw.src[:3], rw.arg)
if s.op is UOps.STORE and len(s.src) == 4 and (rw:=_dfs(s, s.src[3])) != s:
sink_srcs[i] = UOp(rw.op, rw.dtype, rw.src[:3], rw.arg)
sink = UOp(UOps.SINK, None, tuple(sink_srcs))
# do graph rewrite