From 5c43a89fb11bc19d6a5dda724d1920bbe9e56d43 Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Sat, 15 Aug 2026 14:25:31 +0800
Subject: [PATCH] precompile_backward tests for sched_cache (#17544)

* work

* back

* work

* keep +
---
 test/unit/test_schedule_cache.py | 45 ++++++++++++++++++++------------
 1 file changed, 28 insertions(+), 17 deletions(-)

diff --git a/test/unit/test_schedule_cache.py b/test/unit/test_schedule_cache.py
index c304eab636..6b3235903e 100644
--- a/test/unit/test_schedule_cache.py
+++ b/test/unit/test_schedule_cache.py
@@ -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__":