import unittest, pickle, types, tracemalloc import numpy as np from tinygrad import Tensor, Device, TinyJit, Variable, dtypes from tinygrad.helpers import GlobalCounters, ContextVar, Context, DEV from tinygrad.uop.ops import PatternMatcher, UPat, UOp class TestPickle(unittest.TestCase): def test_pickle_code_object(self): y = lambda x: x*2 # noqa: E731 code_str = pickle.dumps(y.__code__) fxn = types.FunctionType(pickle.loads(code_str), globals()) self.assertEqual(fxn(2), 4) def test_pickle_pattern_matcher(self): pm = PatternMatcher([(UPat.cvar('x'), lambda x: x*2)]) sink = UOp.const(2) tt = pm.rewrite(sink) pm_str = pickle.dumps(pm) pm2 = pickle.loads(pm_str) self.assertEqual(pm2.rewrite(sink).key, tt.key) def test_pickle_main_pattern_matcher(self): from tinygrad.uop.symbolic import sym ssym = pickle.dumps(sym) dsym = pickle.loads(ssym) self.assertEqual(dsym.patterns[0][0].location, sym.patterns[0][0].location) def test_pickle_realized_tensor(self): print("** init") t = Tensor.rand(10, 10).realize() st = pickle.dumps(t) t_values = t.numpy() del t # free buffers print("** post pickle") GlobalCounters.reset() t2:Tensor = pickle.loads(st) np.testing.assert_equal(t_values, t2.numpy()) # expect at most one COPY kernel self.assertLessEqual(GlobalCounters.kernel_count, 1) def test_pickle_realized_tensor_alt(self): print("** init") t = Tensor.rand(10, 10).to("CPU").realize() st = pickle.dumps(t) t_values = t.numpy() del t # free buffers print("** post pickle") t2:Tensor = pickle.loads(st) assert t2.uop.is_realized np.testing.assert_equal(t_values, t2.numpy()) def test_pickle_realized_tensor_alt2(self): print("** init") t = Tensor.rand(10, 10).to("CPU").realize() tensor_uop = t.uop assert tensor_uop.is_realized, f"expected {tensor_uop} to be realized" t_values = t.numpy() # pickle st = pickle.dumps(t) # free buffers del t del tensor_uop print("** post pickle") t2:Tensor = pickle.loads(st) assert t2.uop.is_realized, f"expected {t2.uop} to be realized" np.testing.assert_equal(t_values, t2.numpy()) # NOTE: currently Buffer exists on the uop, not tensor def test_pickle_buffer_uop(self): t = Tensor.arange(4).clone().realize() a = t.uop assert a.is_realized self.assertIsNotNone(buffer:=a.base.realized) s = pickle.dumps(a) # free buffers del a del buffer a2:UOp = pickle.loads(s) self.assertListEqual(a2.base.realized.as_memoryview().cast("I").tolist(), [0, 1, 2, 3]) @unittest.skipIf(DEV.interface.startswith("MOCK"), "mock device buffers live in host RAM, not VRAM") def test_pickle_oob_ram(self): N, M = 8, 10**6 ts = [Tensor.rand(M, dtype='float32').realize() for _ in range(N)] tracemalloc.start() st = pickle.dumps(ts, protocol=5, buffer_callback=lambda pb: pb.release()) self.assertLess(tracemalloc.get_traced_memory()[1], N*M*4) tracemalloc.reset_peak() def make_fake_buffers(): for _ in range(N): Device[Device.DEFAULT].synchronize() yield pickle.PickleBuffer(bytearray(M*4)) pickle.loads(st, buffers=make_fake_buffers()) self.assertLess(tracemalloc.get_traced_memory()[1], N*M*4) tracemalloc.stop() def test_pickle_unrealized_tensor(self): t = Tensor.ones(10, 10) st = pickle.dumps(t) t2:Tensor = pickle.loads(st) np.testing.assert_equal(t.numpy(), t2.numpy()) def test_pickle_variable(self): v = Variable("i", 1, 20).bind(10) t1 = Tensor.ones(10, v).contiguous() t2 = Tensor.ones(10, v).contiguous() ret = (t1+t2).sum(1) st = pickle.dumps(ret) del ret vt2 = pickle.loads(st) np.testing.assert_equal(vt2.numpy(), 20) def test_pickle_buffer_view(self): t = Tensor.arange(10).clone(device="CPU").realize() vt = t[3:5].contiguous().realize() assert hasattr(vt.uop.buffer, 'base') ref_value = vt.tolist() st = pickle.dumps(vt) del t, vt vt2 = pickle.loads(st) assert hasattr(vt2.uop.buffer, 'base') assert ref_value == vt2.tolist() def test_pickle_numpy(self): t = Tensor(np.array([1,2,3,4.]), dtype=dtypes.float32) st = pickle.dumps(t) t2:Tensor = pickle.loads(st) np.testing.assert_equal(t.numpy(), t2.numpy()) def test_pickle_jit(self): @TinyJit def add(a, b): return a.sum()+b+1 for _ in range(3): add(Tensor.rand(10, 10), Tensor.rand(10, 10)) st = pickle.dumps(add) del add add_fxn = pickle.loads(st) x = Tensor.ones(10, 10).contiguous().realize() y = Tensor.ones(10, 10).contiguous().realize() print("post jit") out = add_fxn(x, y) np.testing.assert_equal(out.numpy(), 102) def test_pickle_jit_no_del(self): @TinyJit def fn(x): return x + 1.0 for _ in range(3): fn(Tensor.randn(4)) loaded = pickle.loads(pickle.dumps(fn)) self.assertEqual(loaded(Tensor([1.0,2.0,3.0,4.0])).tolist(), [2.0,3.0,4.0,5.0]) def test_pickle_context_var(self): v = ContextVar("test_var", 0) with Context(test_var=1): vs = pickle.dumps(v) v2 = pickle.loads(vs) self.assertEqual(v2.value, 1) def test_pickle_schedule(self): a = Tensor([1,2]) out = a + 2 sched = out.schedule_linear() pk = pickle.dumps(sched) sched_pk = pickle.loads(pk) self.assertEqual(sched_pk.src[-1].src[0], sched.src[-1].src[0]) def test_pickle_renderer(self): from tinygrad.device import Device pk = pickle.dumps(Device.default.renderer) pickle.loads(pk) class TestPickleJIT(unittest.TestCase): @classmethod def setUpClass(cls): N = 10 @TinyJit def add(a, b): return a.sum()+b+1 for _ in range(3): add(Tensor.rand(N, N), Tensor.rand(N, N)) cls.st = pickle.dumps(add) del add def test_inspect(self): import io class FakeClass: def __init__(self, *args, **kwargs): print(self.module, self.name) class InspectUnpickler(pickle.Unpickler): def find_class(self, module, name): return type("SpecializedFakeClass", (FakeClass,), {"name": name, "module": module}) InspectUnpickler(io.BytesIO(self.st)).load() @unittest.skip("we are still saving intermediate buffers") def test_size(self): # confirm no intermediate buffers are saved self.assertLess(len(self.st), 1_000_000) if __name__ == '__main__': unittest.main()