geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

remove some VIEWs we don't need [pr] (#8353)

Browse Source 
* remove some VIEWs we don't need [pr]

* unmasked view and movement op on BUFFER are a part of the spec

---------

Co-authored-by: qazal <[email protected]>
Co-authored-by: qazal <[email protected]>
This commit is contained in:
George Hotz
2024-12-24 14:57:42 +08:00
committed by GitHub
co-authored by qazal qazal
parent 0d6fe6200c
commit b589dec06e
1 changed files with 10 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/engine/schedule.py
+10 -1
View File
@@ -585,7 +585,16 @@ def append_uop(ctx:ScheduleContext, view:UOp, buf_uop:UOp) -> None:
buf_uop.buffer.ref(1)
create_ctx = PatternMatcher([(UPat(Ops.VIEW, name="view", src=(UPat(Ops.BUFFER, name="buf_uop"), UPat())), append_uop)])
remove_movement_ops = PatternMatcher([(UPat(GroupOp.Movement, name="x"), lambda x: x.base.view(unwrap(x.st))),])
remove_movement_ops = PatternMatcher([
(UPat(GroupOp.Movement, name="x"), lambda x: x.base.view(unwrap(x.st))),
# merge one src (unrealized) views
# NOTE: we can't merge realized buffer views here, because the buffer is realized before the view
(UPat(Ops.VIEW, src=(UPat(Ops.VIEW, src=(UPat.var("x"),), name="v1")), name="v2"),
lambda x,v1,v2: v1.replace(arg=v1.arg+v2.arg) if x.op is not Ops.BUFFER else None),
# merge unmasked const views
(UPat(Ops.VIEW, src=(UPat(Ops.VIEW, src=(UPat(), UPat(Ops.CONST)), name="v1")), name="v2"),
lambda v1,v2: v1.replace(arg=v1.arg+v2.arg) if all(v.mask is None for v in v2.st.views) else None),
])
@track_rewrites(named=True)
def create_schedule_with_vars(outs:list[UOp], skip_check:bool=not __debug__) -> tuple[list[ScheduleItem], dict[Variable, int]]: