mirror of
https://github.com/dragonpilot/dragonpilot.git
synced 2026-06-26 16:32:06 +08:00
Vehicle Researcher
db5cbadcf2
date: 2025-06-05T19:54:08 master commit: 8aadf02b2fd91f4e1285e18c2c7feb32d93b66f5
30 lines
818 B
Python
30 lines
818 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
|
|
Device['CPU'].compiler = None # making it not callable
|
|
|
|
try:
|
|
x1 = a1 + b1 + unique_const
|
|
x1.realize() # Same kernel should be from cache.
|
|
finally:
|
|
Device['CPU'].compiler = orig_compile_func
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|