forked from tinygrad/tinygrad
61 lines
2.0 KiB
Python
61 lines
2.0 KiB
Python
import os, unittest
|
|
from unittest.mock import patch
|
|
from tinygrad import Tensor, Variable, Context, UOp
|
|
from tinygrad.callify import transform_to_call
|
|
from tinygrad.helpers import cpu_events
|
|
from tinygrad.schedule import lower_sink_to_linear, schedule_cache
|
|
|
|
def schedule_one():
|
|
(Tensor.empty(1) + 1).schedule_linear()
|
|
|
|
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.linear_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)
|
|
|
|
def test_disk_schedule_cache(self):
|
|
function = transform_to_call(UOp.sink((Tensor.empty(1) + 1).uop))[0].src[0]
|
|
schedule_cache.clear()
|
|
with patch.dict(os.environ, {"DISK_SCACHE":"1"}), \
|
|
patch("tinygrad.schedule.diskcache_get", return_value=None), \
|
|
patch("tinygrad.schedule.diskcache_put") as cache_put:
|
|
lower_sink_to_linear(function)
|
|
cached = cache_put.call_args.args[2]
|
|
|
|
schedule_cache.clear()
|
|
with patch.dict(os.environ, {"DISK_SCACHE":"1"}), \
|
|
patch("tinygrad.schedule.diskcache_get", return_value=cached) as cache_get, \
|
|
patch("tinygrad.schedule.diskcache_put") as cache_put:
|
|
self.assertIs(lower_sink_to_linear(function), cached)
|
|
cache_get.assert_called_once()
|
|
cache_put.assert_not_called()
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|