nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

precompile_backward tests for sched_cache (#17544)

Browse Source 
* work

* back

* work

* keep +
This commit is contained in:
qazal
2026-08-15 15:25:31 +09:00
committed by GitHub
parent e6f5bb9c09
commit 5c43a89fb1
1 changed files with 28 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/unit/test_schedule_cache.py
+28 -17
View File
@@ -4,8 +4,13 @@ from tinygrad import Tensor, Variable, UOp, function
from tinygrad.uop.ops import KernelInfo
from tinygrad.schedule import schedule_cache
def custom_set0_kernel(A:UOp, num:int) -> UOp:
return A[0].set(num).sink(arg=KernelInfo(f"custom_set0_{num}"))
def custom_add_kernel(A:UOp, B:UOp, num:int=0) -> UOp:
return A[0].set(B[0] + num).sink(arg=KernelInfo(f"custom_add_{num}"))
def custom_add_backward(grad_output:UOp, _) -> tuple[None, UOp]:
grad = Tensor.invalids(*grad_output.shape, dtype=grad_output.dtype, device=grad_output.device)
grad = Tensor.custom_kernel(grad, Tensor(grad_output, device=grad_output.device), fxn=functools.partial(custom_add_kernel, num=0))[0]
return None, grad.uop
class TestScheduleCache(unittest.TestCase):
def test_bound_variable_reuses_cache(self):
@@ -25,27 +30,27 @@ class TestScheduleCache(unittest.TestCase):
def test_custom_kernel(self):
for i in range(4):
a = Tensor.empty(1)
a = Tensor.custom_kernel(a, fxn=functools.partial(custom_set0_kernel, num=i))[0]
a, b = Tensor.empty(1), Tensor.ones(1)
a = Tensor.custom_kernel(a, b, fxn=functools.partial(custom_add_kernel, num=i))[0]
a.realize()
self.assertEqual(a.item(), i)
self.assertEqual(a.item(), i+1)
def test_same_custom_function_reuses_cache(self):
schedule_cache.clear()
fxn = functools.partial(custom_set0_kernel, num=10)
fxn = functools.partial(custom_add_kernel, num=10)
# first run
a = Tensor.empty(1)
a = Tensor.custom_kernel(a, fxn=fxn)[0]
a, x = Tensor.empty(1), Tensor.ones(1)
a = Tensor.custom_kernel(a, x, fxn=fxn)[0]
a.realize()
self.assertEqual(a.item(), 10)
self.assertEqual(a.item(), 11)
cache_size_after_first = len(schedule_cache)
# second run with same function should reuse cache
b = Tensor.empty(1)
b = Tensor.custom_kernel(b, fxn=fxn)[0]
b, x = Tensor.empty(1), Tensor.ones(1)
b = Tensor.custom_kernel(b, x, fxn=fxn)[0]
b.realize()
self.assertEqual(b.item(), 10)
self.assertEqual(b.item(), 11)
self.assertEqual(len(schedule_cache), cache_size_after_first)
def test_simple(self):
@@ -67,21 +72,27 @@ class TestScheduleCache(unittest.TestCase):
@unittest.expectedFailure
def test_simple_precompile(self):
@function(precompile=True)
@function(precompile=True, precompile_backward=True)
def f(x:Tensor) -> Tensor:
out = Tensor.invalids(*x.shape, dtype=x.dtype, device=x.device)
out = Tensor.custom_kernel(out, fxn=functools.partial(custom_set0_kernel, num=10))[0]
out = Tensor.custom_kernel(out, x, fxn=functools.partial(custom_add_kernel, num=10), grad_fxn=custom_add_backward)[0]
return out + x
# warmup
x = Tensor.ones(1).realize()
_ = f(x).realize()
out = f(x)
out.backward(x)
self.assertEqual(out.item(), 12)
self.assertEqual(x.grad.item(), 2)
# use the cache next time function is called
start_len_schedule_cache = len(schedule_cache)
for _ in range(3):
num = f(x).realize()
self.assertEqual(num.item(), 11)
x = Tensor.ones(1).realize()
out = f(x)
out.backward(x)
self.assertEqual(out.item(), 12)
self.assertEqual(x.grad.item(), 2)
self.assertEqual(len(schedule_cache), start_len_schedule_cache)
if __name__ == "__main__":