github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-09 15:08:02 -05:00
Code Issues Packages Projects Releases Wiki Activity

move all to compile api (#2203)

Browse Source 
* move metal+clang to compile api

* all to the new style

* remove binary arg

* fix triton

* fixup tests

* fix clang

* diskcache is generic

* __wrapped__

* compile_gpu

* fix thneed

* keep the src in the ASTRunner

* lib

* move compile_gpu

* compile_gpu in device

* put compiler in astrunner

* test reverts

* triton compiler

* ugh, that too
This commit is contained in:
George Hotz
2023-11-01 23:01:32 -07:00
committed by GitHub
parent 8932816816
commit 03cf0afa4f
18 changed files with 128 additions and 136 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
test/test_kernel_cache.py
View File

@@ -2,36 +2,34 @@
import unittest
import secrets
import string
import tempfile
import pathlib
from tinygrad.tensor import Tensor
from tinygrad.ops import Device
from tinygrad.helpers import cache_compiled
import tinygrad.runtime.ops_clang
from tinygrad.helpers import diskcache
def generate_random_string(length=16):
alphabet = string.ascii_letters + string.digits
return ''.join(secrets.choice(alphabet) for _ in range(length))
compile_call_count = 0
@diskcache
def helper_test_compile(prg:str) -> bytes:
global compile_call_count
compile_call_count += 1
return prg.encode()
class TestKernelCache(unittest.TestCase):
compile_call_count = 0
@cache_compiled
def __helper_test_compile(self, prg, output_file=pathlib.Path(tempfile.mktemp()), **kwargs):
self.compile_call_count += 1
return prg.encode()
def test_compile_cache(self):
prg1 = generate_random_string(64) + "a"
prg2 = generate_random_string(64) + "b"
cold_compile_res = self.__helper_test_compile(prg1)
warm_compile_res = self.__helper_test_compile(prg1)
cold_compile_res = helper_test_compile(prg1)
warm_compile_res = helper_test_compile(prg1)
assert cold_compile_res == warm_compile_res == prg1.encode()
assert self.compile_call_count == 1
assert compile_call_count == 1
prg2_res = self.__helper_test_compile(prg2)
prg2_res = helper_test_compile(prg2)
assert prg2_res == prg2.encode()
assert self.compile_call_count == 2
assert compile_call_count == 2
def test_kernel_cache_in_action(self):
if Device.DEFAULT not in ["CLANG"]:
@@ -42,15 +40,15 @@ class TestKernelCache(unittest.TestCase):
x = a + b
x.realize()
orig_compile_func = tinygrad.runtime.ops_clang.ClangBuffer.runtime.compile
tinygrad.runtime.ops_clang.ClangBuffer.runtime.compile = None # making it not callable
orig_compile_func = Device['CLANG'].compiler
Device['CLANG'].compiler = None # making it not callable
a1 = Tensor.rand(4,4)
b1 = Tensor.rand(4,4)
x1 = a1 + b1
x1.realize() # Same kernel should be from cache.
tinygrad.runtime.ops_clang.ClangBuffer.runtime.compile = orig_compile_func
Device['CLANG'].compiler = orig_compile_func
if __name__ == "__main__":
unittest.main()