mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-04-29 03:00:14 -04:00
42 lines
1.1 KiB
Python
42 lines
1.1 KiB
Python
import unittest
|
|
from tinygrad import Tensor, Variable, Context
|
|
from tinygrad.helpers import cpu_events
|
|
from tinygrad.schedule import schedule_cache
|
|
|
|
def schedule_one():
|
|
Tensor([1]).schedule()
|
|
|
|
class TestScheduleCache(unittest.TestCase):
|
|
def test_bound_variable_var_vals(self):
|
|
v = Variable('pos', 1, 100)
|
|
x = Tensor.ones(10).contiguous().realize()
|
|
|
|
t = x + Tensor(v.bind(42))
|
|
_, var_vals = t.schedule_with_vars()
|
|
self.assertEqual(var_vals, {'pos': 42})
|
|
|
|
def test_disable_schedule_cache(self):
|
|
schedule_cache.clear()
|
|
|
|
# test write
|
|
with Context(SCACHE=0): schedule_one()
|
|
self.assertEqual(len(schedule_cache), 0)
|
|
with Context(SCACHE=1):
|
|
schedule_one()
|
|
schedule_one()
|
|
self.assertEqual(len(schedule_cache), 1)
|
|
|
|
# test read
|
|
with Context(PROFILE=1):
|
|
cpu_events.clear()
|
|
with Context(SCACHE=0): schedule_one()
|
|
num_events_no_cache = len(cpu_events)
|
|
|
|
cpu_events.clear()
|
|
with Context(SCACHE=1): schedule_one()
|
|
num_events_cache = len(cpu_events)
|
|
self.assertLess(num_events_cache, num_events_no_cache)
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|