mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-01-09 15:08:02 -05:00
c5bd28e21d
* start work on schedule cache * local unique * schedule cache works * schedule cache cleanup * fix tests * preserve metadata * oops, fix cache * put that there * fix spec * always miss * why is that broken? * src[0].op * fix process replay * delete abstractions2 * reenable the actual schedule cache * metadata is best effort * fix JIT in examples/gradaccum_mnist.py * full jit * fixed and test is real
87 lines
3.0 KiB
Python
87 lines
3.0 KiB
Python
import gc
|
|
from tinygrad import Tensor, UOp, Device, nn
|
|
from tinygrad.engine.schedule import schedule_cache
|
|
from tinygrad.engine.realize import method_cache, get_program
|
|
from tinygrad.schedule.indexing import apply_movement_op, _apply_reshape
|
|
from tinygrad.uop.divandmod import fold_divmod_general
|
|
from test.test_tiny import TestTiny
|
|
|
|
def uops_allocated(): return sum([isinstance(x, UOp) for x in gc.get_objects()])
|
|
def print_uops():
|
|
for x in gc.get_objects():
|
|
if isinstance(x, UOp): print(x)
|
|
|
|
def start(): pass
|
|
def single_tensor(): Tensor([2])
|
|
def two_plus_two(): Tensor([2])+Tensor([2])
|
|
def two_plus_two_schedule(): (Tensor([2])+Tensor([2])).schedule()
|
|
def two_plus_two_kernel():
|
|
si = (Tensor([2])+Tensor([2])).schedule()[-1]
|
|
get_program(si.ast, Device.default.renderer)
|
|
def two_plus_two_linearize():
|
|
si = (Tensor([2])+Tensor([2])).schedule()[-1]
|
|
get_program(si.ast, Device.default.renderer)
|
|
def two_plus_two_realize(): (Tensor([2])+Tensor([2])).realize()
|
|
def two_plus_two_item(): (Tensor([2])+Tensor([2])).item()
|
|
def gradient_test():
|
|
x = Tensor.eye(3, requires_grad=True)
|
|
y = Tensor([[2.0,0,-2.0]], requires_grad=True)
|
|
z = y.matmul(x).sum()
|
|
z.backward()
|
|
def realized_eye():
|
|
Tensor.eye(3, requires_grad=True).realize()
|
|
def realized_list():
|
|
Tensor([[2.0,0,-2.0]], requires_grad=True).realize()
|
|
def kernel_matmul():
|
|
x = Tensor.eye(3, requires_grad=True)
|
|
y = Tensor([[2.0,0,-2.0]], requires_grad=True)
|
|
z = y.matmul(x)
|
|
si = z.schedule()[-1]
|
|
get_program(si.ast, Device.default.renderer)
|
|
def realized_matmul():
|
|
x = Tensor.eye(3, requires_grad=True)
|
|
y = Tensor([[2.0,0,-2.0]], requires_grad=True)
|
|
z = y.matmul(x)
|
|
Tensor.realize(z)
|
|
def realized_gradient():
|
|
x = Tensor.eye(3, requires_grad=True)
|
|
y = Tensor([[2.0,0,-2.0]], requires_grad=True)
|
|
z = y.matmul(x).sum()
|
|
z.backward()
|
|
Tensor.realize(x, y, z, x.grad, y.grad)
|
|
def nn_batchnorm(): nn.BatchNorm(64)
|
|
def nn_conv2d(): nn.Conv2d(64, 64, 3)
|
|
def plus(): TestTiny().test_plus()
|
|
def mnist(): TestTiny().test_mnist()
|
|
def mnist_backward(): TestTiny().test_mnist_backward()
|
|
|
|
tests = [start, single_tensor, two_plus_two, two_plus_two_schedule, two_plus_two_kernel,
|
|
two_plus_two_linearize, two_plus_two_realize, two_plus_two_item, gradient_test,
|
|
realized_eye, realized_list, kernel_matmul, realized_matmul, realized_gradient,
|
|
nn_batchnorm, nn_conv2d, plus, mnist, mnist_backward]
|
|
|
|
if __name__ == "__main__":
|
|
gc.disable()
|
|
start_uops = uops_allocated()
|
|
# there's a few consts created as default values
|
|
print_uops()
|
|
for t in tests:
|
|
t()
|
|
|
|
# these caches will keep uops alive
|
|
schedule_cache.clear()
|
|
method_cache.clear()
|
|
apply_movement_op.cache_clear()
|
|
_apply_reshape.cache_clear()
|
|
fold_divmod_general.cache_clear()
|
|
UOp.const.cache_clear()
|
|
Tensor._device_seeds.clear()
|
|
Tensor._device_rng_counters.clear()
|
|
|
|
new_uops = uops_allocated()
|
|
gc.collect()
|
|
new_uops_gc = uops_allocated()
|
|
print(f"{t.__name__:30s}: {new_uops:3d} -> {new_uops_gc:3d}")
|
|
if new_uops != start_uops: print_uops()
|
|
assert new_uops == start_uops
|