github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-29 03:00:14 -04:00
Code Issues Packages Projects Releases Wiki Activity

remove movementops in viz graph rewrite [pr] (#8665)

Browse Source
This commit is contained in:
qazal
2025-01-18 06:34:12 -05:00
committed by GitHub
parent b49a04145e
commit 5267a411e7
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tinygrad/engine/schedule.py
View File

@@ -520,8 +520,8 @@ remove_movement_ops = PatternMatcher([
@track_rewrites(named=True)
def create_schedule_with_vars(outs:list[UOp], skip_check:bool=not __debug__) -> tuple[list[ScheduleItem], dict[Variable, int], dict[UOp, UOp]]:
big_sink = UOp.sink(*outs)
# if using VIZ, do an empty graph rewrite to vizualize the Tensor graph
if getenv("VIZ"): graph_rewrite(big_sink, PatternMatcher([]))
# if using VIZ, do a graph rewrite to vizualize the Tensor graph
if getenv("VIZ"): graph_rewrite(big_sink, remove_movement_ops)
if not skip_check: type_verify(list(big_sink.toposort), tensor_uop_spec)
# to_uop is removing (many) of the movement ops
sink = add_buffers(big_sink, ctx:=ScheduleContext(), cache={})