tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 13:56:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

one less rewrite on schedule [pr] (#10872)

Browse Source 
* one less rewrite on schedule [pr]

* verify in ebs
This commit is contained in:
George Hotz
2025-06-18 17:06:17 -07:00
committed by GitHub
parent e7a26211d2
commit 18593c9800
2 changed files with 5 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/external/external_benchmark_schedule.py
Vendored
+3
View File
@@ -7,6 +7,7 @@ from tinygrad.codegen.kernel import Kernel
from tinygrad.codegen.heuristic import hand_coded_optimizations
from tinygrad.codegen import get_rewrites_for_renderer, apply_rewrites, rewrites_for_linearizer
from tinygrad.engine.search import beam_search, bufs_from_lin
from tinygrad.uop.spec import type_verify
if __name__ == "__main__":
mdl = ResNet50()
@@ -56,4 +57,6 @@ if __name__ == "__main__":
uops_line = []
for u in rewritten_uops:
uops_line.append(apply_rewrites(u, rewrites_for_linearizer))
with Timing("***** model verify in "):
for u in uops_line: type_verify(u.arg.lst)
print(sum(len(u.arg.lst) for u in uops_line))
tinygrad/engine/kernelize.py
+2 -4
View File
@@ -430,10 +430,8 @@ def get_kernelize_map(big_sink:UOp) -> dict[UOp, UOp]:
# insert gbarriers in places determined by the realize map
realize_map = group_realizes(tensor_map[big_sink])
tensor_map = graph_rewrite_map(tensor_map[big_sink], add_gbarrier, realize_map, bottom_up=True, input_map=tensor_map, name="insert_gbarrier")
# optionally reorder gbarriers or insert more (top down)
tensor_map = graph_rewrite_map(tensor_map[big_sink], finalize_gbarrier, input_map=tensor_map, name="finalize_gbarrier")
tensor_map = graph_rewrite_map(tensor_map[big_sink], remove_tags, input_map=tensor_map, name="remove_tags")
tensor_map = graph_rewrite_map(tensor_map[big_sink], add_gbarrier, ctx=realize_map, bottom_up=True, input_map=tensor_map, name="insert_gbarrier")
tensor_map = graph_rewrite_map(tensor_map[big_sink], finalize_gbarrier+remove_tags, input_map=tensor_map, name="finalize_gbarrier")
# TODO: move view_left/view_right here