move all to compile api (#2203)

* 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

test/external/external_test_speed_llama.py

@@ -11,7 +11,7 @@ from tinygrad.helpers import dtypes, prod
 from tinygrad.runtime.lib import RawBuffer
 
 class FakeProgram:
-  def __init__(self, name:str, prg:str, binary:bool): pass
+  def __init__(self, name:str, prg:str): pass
   def __call__(self, global_size, local_size, *bufs, wait=False): pass
 
 class RawFakeBuffer(RawBuffer):

test/test_custom_function.py

@@ -24,7 +24,7 @@ def atan2_gpu(ret:LazyBuffer, a:LazyBuffer, b:LazyBuffer):
     __kernel void atan2_gpu(global float *c, global float *a, global float *b) {
       int idx = get_global_id(0);
       c[idx] = atan2(a[idx], b[idx]);
-    }""", global_size=[prod(ret.shape)]).build(Device[ret.device].runtime).exec([ret.realized, a.realized, b.realized])
+    }""", global_size=[prod(ret.shape)]).build(Device[ret.device].compiler, Device[ret.device].runtime).exec([ret.realized, a.realized, b.realized])
   return ret.realized
 
 def atan2_cpu(ret:LazyBuffer, a:LazyBuffer, b:LazyBuffer):

test/test_kernel_cache.py

@@ -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()

test/test_uops.py

@@ -8,7 +8,7 @@ from tinygrad.codegen.linearizer import UOps, UOp
 
 def _uops_to_prg(uops):
   src, runtime_args = Device[Device.DEFAULT].renderer("test", uops)
-  return ASTRunner("test", src, [1], [1], runtime_args=runtime_args).build(Device[Device.DEFAULT].runtime)
+  return ASTRunner("test", src, [1], [1], runtime_args=runtime_args).build(Device[Device.DEFAULT].compiler, Device[Device.DEFAULT].runtime)
 
 def uop(uops:List[UOp], uop:UOps, dtype:Optional[DType], vin:Tuple[UOp, ...], arg:Any=None) -> UOp:
   uops.append(UOp(uop, dtype, tuple(vin), arg, len(uops)))