Files
StarPilot/tinygrad_repo/test/test_kernel_cache.py
T
github-actions[bot] eefc084302 sunnypilot v2025.12.18-3820
version: sunnypilot v2025.003.000 (dev)
date: 2025-12-18T05:48:43
master commit: 16c052af0835f049a67b80251d8cc1114fc08bb4
2025-12-18 05:48:43 +00:00

30 lines
863 B
Python

#!/usr/bin/env python
import unittest
from tinygrad.tensor import Tensor
from tinygrad import Device
class TestKernelCache(unittest.TestCase):
def test_kernel_cache_in_action(self):
if Device.DEFAULT not in ["CPU"]:
self.skipTest("No custom kernel cache is implemented")
unique_const = 0.6765677269
a = Tensor.rand(4,4).realize()
b = Tensor.rand(4,4).realize()
x = a + b + unique_const
x.realize()
a1 = Tensor.rand(4,4).realize()
b1 = Tensor.rand(4,4).realize()
orig_compile_func = Device['CPU'].compiler.compile_cached
Device['CPU'].compiler.compile_cached = None # making it not callable
try:
x1 = a1 + b1 + unique_const
x1.realize() # Same kernel should be from cache.
finally:
Device['CPU'].compiler.compile_cached = orig_compile_func
if __name__ == "__main__":
unittest.main()