github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-10 07:28:15 -05:00
Code Issues Packages Projects Releases Wiki Activity

remove ScheduleItem and merge it with ExecItem (#13759)

Browse Source 
* remove ExecItem and merge it with ScheduleItem

* less diff

* fix issues

* min diff

* don't change bufs in _lower

* min diff

* update

* revert

* fixes

* diff
This commit is contained in:
George Hotz
2025-12-19 17:04:24 -04:00
committed by GitHub
parent df6cde8a00
commit 744af193f0
35 changed files with 172 additions and 395 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
test/testextra/test_tk.py
View File

@@ -2,7 +2,8 @@ import unittest, math, time
from tinygrad import Tensor, Device, dtypes, Context
from tinygrad.uop.ops import UOp, Ops
from tinygrad.engine.realize import ExecItem, get_runner
from tinygrad.engine.realize import get_runner
from tinygrad.engine.schedule import ExecItem
from tinygrad.engine.jit import TinyJit
from tinygrad.helpers import CI
import numpy as np
@@ -64,7 +65,7 @@ class TestTK(unittest.TestCase):
c = Tensor.empty(1, 1, N, N, dtype="float32")
Tensor.realize(a, b, c)
ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (c, a, b)])
ei = ExecItem(sink, [t.uop.buffer for t in (c, a, b)], prg=get_runner(Device.DEFAULT, sink))
for _ in range(5): ei.run(wait=True)
c = c.float()
@@ -113,7 +114,7 @@ class TestTK(unittest.TestCase):
c = Tensor.empty(1, 1, N, N, dtype="float32")
Tensor.realize(a, b, c)
ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (c, a, b)])
ei = ExecItem(sink, [t.uop.buffer for t in (c, a, b)], prg=get_runner(Device.DEFAULT, sink))
for _ in range(5): ei.run(wait=True)
c = c.float()
@@ -149,7 +150,7 @@ class TestTK(unittest.TestCase):
b = Tensor.empty(1, 1, N, N, dtype="float32")
Tensor.realize(a, b)
ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
ei = ExecItem(sink, [t.uop.buffer for t in (b, a)], prg=get_runner(Device.DEFAULT, sink))
for _ in range(5): ei.run(wait=True)
b = b.float()
@@ -188,7 +189,7 @@ class TestTK(unittest.TestCase):
b = Tensor.empty(1, 1, N, N, dtype="float32")
Tensor.realize(a, b)
ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
ei = ExecItem(sink, [t.uop.buffer for t in (b, a)], prg=get_runner(Device.DEFAULT, sink))
for _ in range(5): ei.run(wait=True)
b = b.float()
@@ -230,7 +231,7 @@ class TestTK(unittest.TestCase):
c = Tensor.empty(1, 1, N, N, dtype="float32")
Tensor.realize(a, b, c)
ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, c, a)])
ei = ExecItem(sink, [t.uop.buffer for t in (b, c, a)], prg=get_runner(Device.DEFAULT, sink))
for _ in range(5): ei.run(wait=True)
b = b.float()
c = c.float()
@@ -270,7 +271,7 @@ class TestTK(unittest.TestCase):
b = Tensor.empty(1, 1, N, N, dtype="float32")
Tensor.realize(a, b)
ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
ei = ExecItem(sink, [t.uop.buffer for t in (b, a)], prg=get_runner(Device.DEFAULT, sink))
for _ in range(5): ei.run(wait=True)
b = b.float()
@@ -307,7 +308,7 @@ class TestTK(unittest.TestCase):
b = Tensor.empty(1, 1, N, N, dtype="float32")
Tensor.realize(a, b)
ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
ei = ExecItem(sink, [t.uop.buffer for t in (b, a)], prg=get_runner(Device.DEFAULT, sink))
for _ in range(5): ei.run(wait=True)
b = b.float()
@@ -352,7 +353,7 @@ class TestTK(unittest.TestCase):
b = Tensor.empty(1, 1, N, N, dtype="float32")
Tensor.realize(a, b)
ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
ei = ExecItem(sink, [t.uop.buffer for t in (b, a)], prg=get_runner(Device.DEFAULT, sink))
for _ in range(5): ei.run(wait=True)
b = b.float()
@@ -397,7 +398,7 @@ class TestTK(unittest.TestCase):
b = Tensor.empty(1, 1, N, M, dtype="float32")
Tensor.realize(a, b)
ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
ei = ExecItem(sink, [t.uop.buffer for t in (b, a)], prg=get_runner(Device.DEFAULT, sink))
for _ in range(5): ei.run(wait=True)
b = b.float()
@@ -442,7 +443,7 @@ class TestTK(unittest.TestCase):
b = Tensor.empty(1, 1, N, N, dtype="float32")
Tensor.realize(a, b)
ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
ei = ExecItem(sink, [t.uop.buffer for t in (b, a)], prg=get_runner(Device.DEFAULT, sink))
for _ in range(5): ei.run(wait=True)
b = b.float()
@@ -487,7 +488,7 @@ class TestTK(unittest.TestCase):
b = Tensor.empty(1, 1, N, M, dtype="float32")
Tensor.realize(a, b)
ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
ei = ExecItem(sink, [t.uop.buffer for t in (b, a)], prg=get_runner(Device.DEFAULT, sink))
for _ in range(5): ei.run(wait=True)
b = b.float()
@@ -547,7 +548,7 @@ class TestTK(unittest.TestCase):
b = Tensor.empty(1, 1, BLOCK_SIZE, N, dtype="float32")
Tensor.realize(a, b)
ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
ei = ExecItem(sink, [t.uop.buffer for t in (b, a)], prg=get_runner(Device.DEFAULT, sink))
for _ in range(5): ei.run(wait=True)
b = b.float()
@@ -607,7 +608,7 @@ class TestTK(unittest.TestCase):
b = Tensor.empty(1, 1, N, BLOCK_SIZE, dtype="float32")
Tensor.realize(a, b)
ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
ei = ExecItem(sink, [t.uop.buffer for t in (b, a)], prg=get_runner(Device.DEFAULT, sink))
for _ in range(5): ei.run(wait=True)
b = b.float()
@@ -717,7 +718,7 @@ class TestTK(unittest.TestCase):
out = Tensor.empty(B, N, H, D, dtype=dtypes.bfloat16)
Tensor.realize(q, k, v, out)
ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (out, q, k, v)])
ei = ExecItem(sink, [t.uop.buffer for t in (out, q, k, v)], prg=get_runner(Device.DEFAULT, sink))
for _ in range(5):
et = ei.run(wait=True)
attn_flops = 2 * B * H * N * N * D + \