forked from tinygrad/tinygrad
don't create linearizer if we are in the method cache (#1969)
* don't create linearizer if we are in the method cache * remove unchecked properties * that key isn't used * fix default type is sticky
This commit is contained in:
George Hotz
Vendored
-48
@@ -1,48 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
import unittest
|
||||
import numpy as np
|
||||
from tinygrad.ops import LazyOp, BinaryOps, ReduceOps
|
||||
from tinygrad.runtime.ops_llvm import LLVMBuffer
|
||||
|
||||
class TestLLVM(unittest.TestCase):
|
||||
def test_add(self):
|
||||
a = LLVMBuffer.fromCPU(np.ones((4,4)))
|
||||
b = LLVMBuffer.fromCPU(np.ones((4,4)))
|
||||
ast = LazyOp(BinaryOps.ADD, (a,b))
|
||||
ret = LLVMBuffer((4,4)).exec_ast(ast)
|
||||
print(ret.toCPU())
|
||||
|
||||
def test_sum(self):
|
||||
a = LLVMBuffer.fromCPU(np.ones((4,4)))
|
||||
ast = LazyOp(ReduceOps.SUM, (a,), (1,1))
|
||||
ret = LLVMBuffer((1,1)).exec_ast(ast)
|
||||
print(ret.toCPU())
|
||||
|
||||
def test_sum_add(self):
|
||||
a = LLVMBuffer.fromCPU(np.ones((4,4)))
|
||||
b = LLVMBuffer.fromCPU(np.ones((1,1)))
|
||||
ast = LazyOp(ReduceOps.SUM, (a,), (1,1))
|
||||
ast = LazyOp(BinaryOps.ADD, (ast,b))
|
||||
ret = LLVMBuffer((1,1)).exec_ast(ast)
|
||||
print(ret.toCPU())
|
||||
|
||||
def test_add_sum(self):
|
||||
a = LLVMBuffer.fromCPU(np.ones((4,4)))
|
||||
b = LLVMBuffer.fromCPU(np.ones((4,4)))
|
||||
ast = LazyOp(BinaryOps.ADD, (a,b))
|
||||
ast = LazyOp(ReduceOps.SUM, (ast,), (1,1))
|
||||
ret = LLVMBuffer((1,1)).exec_ast(ast)
|
||||
print(ret.toCPU())
|
||||
|
||||
def test_add_sum_add(self):
|
||||
a = LLVMBuffer.fromCPU(np.ones((4,4)))
|
||||
b = LLVMBuffer.fromCPU(np.ones((4,4)))
|
||||
c = LLVMBuffer.fromCPU(np.ones((1,1)))
|
||||
ast = LazyOp(BinaryOps.ADD, (a,b))
|
||||
ast = LazyOp(ReduceOps.SUM, (ast,), (1,1))
|
||||
ast = LazyOp(BinaryOps.ADD, (ast,c))
|
||||
ret = LLVMBuffer((1,1)).exec_ast(ast)
|
||||
print(ret.toCPU())
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -67,12 +67,14 @@ def derandomize_model(model):
|
||||
|
||||
class TestRealWorld(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.old_type = Tensor.default_type
|
||||
np.random.seed(2002)
|
||||
if getenv("KOPT"):
|
||||
self.oldfunc = getattr(__import__("tinygrad.codegen.search", fromlist=["kernel_optimize_search"]), "kernel_optimize_search")
|
||||
setattr(__import__("tinygrad.codegen.search", fromlist=["kernel_optimize_search"]), "kernel_optimize_search", kopt_search_hook)
|
||||
|
||||
def tearDown(self):
|
||||
Tensor.default_type = self.old_type
|
||||
if getenv("KOPT"):
|
||||
setattr(__import__("tinygrad.codegen.search", fromlist=["kernel_optimize_search"]), "kernel_optimize_search", self.oldfunc)
|
||||
|
||||
@@ -86,7 +88,6 @@ class TestRealWorld(unittest.TestCase):
|
||||
|
||||
@unittest.skipUnless(Device.DEFAULT in JIT_SUPPORTED_DEVICE and Device.DEFAULT not in ["LLVM"], "needs JIT, too long on CI LLVM")
|
||||
def test_llama(self):
|
||||
old_type = Tensor.default_type
|
||||
Tensor.default_type = dtypes.float16
|
||||
|
||||
args_tiny = {"dim": 1024, "multiple_of": 256, "n_heads": 8, "n_layers": 8, "norm_eps": 1e-05, "vocab_size": 1000}
|
||||
@@ -97,11 +98,8 @@ class TestRealWorld(unittest.TestCase):
|
||||
# NOTE: only test one pass, not testing the dynamic shape autoregressive part
|
||||
helper_test("test_llama", lambda: (Tensor([[1,]]),), test, 0.22 if CI else 13.5, 126 if CI else 486)
|
||||
|
||||
Tensor.default_type = old_type
|
||||
|
||||
@unittest.skipUnless(Device.DEFAULT in JIT_SUPPORTED_DEVICE and Device.DEFAULT not in ["LLVM"], "needs JIT, too long on CI LLVM")
|
||||
def test_gpt2(self):
|
||||
old_type = Tensor.default_type
|
||||
Tensor.default_type = dtypes.float16
|
||||
|
||||
args_tiny = {"dim": 1024, "n_heads": 8, "n_layers": 8, "norm_eps": 1e-5, "vocab_size": 1000}
|
||||
@@ -111,8 +109,6 @@ class TestRealWorld(unittest.TestCase):
|
||||
def test(t): return model(t, 0).realize()
|
||||
helper_test("test_gpt2", lambda: (Tensor([[1,]]),), test, 0.21 if CI else 0.9, 129 if CI else 369)
|
||||
|
||||
Tensor.default_type = old_type
|
||||
|
||||
@unittest.skipIf(getenv("KOPT"), "cifar hangs with KOPT")
|
||||
@unittest.skipUnless(Device.DEFAULT in JIT_SUPPORTED_DEVICE and Device.DEFAULT not in ["LLVM"], "needs JIT, too long on CI LLVM")
|
||||
def test_train_cifar(self):
|
||||
|
||||
@@ -53,7 +53,6 @@ class Kernel:
|
||||
self.opts = opts if opts else (cast(Compiled, Device[Device.DEFAULT]).linearizer_opts if isinstance(Device[Device.DEFAULT], Compiled) else LinearizerOptions())
|
||||
self.ast = ast
|
||||
self.var_vals = var_vals
|
||||
self.key = (ast, tuple(var_vals.keys())) if var_vals else ast
|
||||
|
||||
# fetch lazyop info
|
||||
self.info: FlopCounter = get_lazyop_info(cast(LazyOp, self.ast))
|
||||
|
||||
@@ -50,10 +50,10 @@ def kernel_optimize_search(k:Linearizer, create_k:Callable[[], Linearizer], to_p
|
||||
|
||||
# optimization
|
||||
global_db = None
|
||||
def kernel_optimize(k:Linearizer, create_k:Callable[[], Linearizer], to_prg, bufs):
|
||||
def kernel_optimize(k:Linearizer, create_k:Callable[[], Linearizer], to_prg, bufs, key):
|
||||
global global_db
|
||||
|
||||
skey = str(k.key)
|
||||
skey = str(key)
|
||||
|
||||
if getenv("KOPT") == 2 and global_db is None:
|
||||
import shelve
|
||||
|
||||
+13
-15
@@ -67,14 +67,9 @@ class LazyOp:
|
||||
@property
|
||||
def st(self): raise NotImplementedError
|
||||
@property
|
||||
def children(self): raise NotImplementedError
|
||||
@property
|
||||
def shape(self): raise NotImplementedError
|
||||
@property
|
||||
def realized(self): raise NotImplementedError
|
||||
@property
|
||||
def optype(self): raise NotImplementedError
|
||||
def realize(self): raise NotImplementedError
|
||||
def children(self): raise NotImplementedError
|
||||
|
||||
# movement ops
|
||||
def reshape(self, _): raise NotImplementedError
|
||||
@@ -117,7 +112,7 @@ class Interpreted:
|
||||
if ast.op in BufferOps and ast.op not in self.fxn_for_op:
|
||||
if ast.op == BufferOps.MEM:
|
||||
assert inputs[ast.arg.idx-1].dtype == ast.arg.dtype, "dtype mismatch"
|
||||
buf = self.to_underlying(inputs[ast.arg.idx-1])
|
||||
buf = self.to_underlying(inputs[ast.arg.idx-1].realized)
|
||||
elif ast.op == BufferOps.CONST:
|
||||
buf = self.to_underlying(self.buffer.fromCPU(np.array(ast.arg.val, dtype=ast.arg.dtype.np)))
|
||||
for mop,arg in ast.arg.st.to_movement_ops(): buf = self.fxn_for_op[mop](buf, arg)
|
||||
@@ -226,7 +221,7 @@ class Compiled:
|
||||
if output.realized:
|
||||
for i,a in enumerate(inputs):
|
||||
# TODO: if this is contiguous it's fine
|
||||
if a == output.realized:
|
||||
if a.realized == output.realized:
|
||||
if any(not x.arg.st.contiguous for x in ast.get_lazyops() if x.op == BufferOps.MEM and x.arg.idx == i+1):
|
||||
output.realized = None
|
||||
break
|
||||
@@ -238,23 +233,26 @@ class Compiled:
|
||||
from tinygrad.jit import CacheCollector
|
||||
CacheCollector._mark_output_buffer(output.output_buffer)
|
||||
|
||||
from tinygrad.codegen.linearizer import Linearizer
|
||||
k = Linearizer(ast, self.linearizer_opts, var_vals)
|
||||
# all the rawbuffers
|
||||
rawbuffers = [output.realized] + [x.realized for x in inputs]
|
||||
key = (ast, tuple(var_vals.keys())) if var_vals else ast # TODO: remove var_vals so the key can just be the AST
|
||||
|
||||
# compilation time
|
||||
def get_program():
|
||||
from tinygrad.codegen.linearizer import Linearizer
|
||||
k = Linearizer(ast, self.linearizer_opts, var_vals)
|
||||
from tinygrad.codegen.search import kernel_optimize
|
||||
if getenv("KOPT"): kernel_optimize(k, lambda: Linearizer(ast, self.linearizer_opts, var_vals), self.to_program, [output.realized]+inputs)
|
||||
if getenv("KOPT"): kernel_optimize(k, lambda: Linearizer(ast, self.linearizer_opts, var_vals), self.to_program, rawbuffers, key)
|
||||
elif not getenv("NOOPT"): k.hand_coded_optimizations()
|
||||
return self.to_program(k)
|
||||
|
||||
if hasattr(k, 'key') and getenv("ENABLE_METHOD_CACHE", 1):
|
||||
if k.key not in self.method_cache: self.method_cache[k.key] = get_program()
|
||||
prg = self.method_cache[k.key]
|
||||
if getenv("ENABLE_METHOD_CACHE", 1):
|
||||
if key not in self.method_cache: self.method_cache[key] = get_program()
|
||||
prg = self.method_cache[key]
|
||||
else:
|
||||
prg = get_program()
|
||||
|
||||
if prg.name == getenv("PRINT_PRG", ''): print(prg.prg)
|
||||
|
||||
prg.exec([output.realized]+inputs, var_vals=var_vals)
|
||||
prg.exec(rawbuffers, var_vals=var_vals)
|
||||
return output.realized
|
||||
|
||||
+1
-1
@@ -23,7 +23,7 @@ def run_schedule(schedule:List[Tuple[LazyOp, LazyBuffer, Tuple[LazyBuffer, ...]]
|
||||
for i,s in enumerate(op.src): assert isinstance(s, LazyOp) and s.op == BufferOps.MEM and s.arg.idx == i+1 and s.arg.st.contiguous, f"bad LoadOps src {i}: {s}"
|
||||
LOAD_OPS_DISPATCHER[cast(LoadOps, op.op)](out, *buffers)
|
||||
else:
|
||||
out.realized = Device[out.device].exec_ast(op, output=out, inputs=[x.realized for x in buffers], var_vals=out.var_vals, **out._device_extra_args())
|
||||
out.realized = Device[out.device].exec_ast(op, output=out, inputs=buffers, var_vals=out.var_vals, **out._device_extra_args())
|
||||
del out.op
|
||||
for v in out.views: del v.op
|
||||
assert out.realized and isinstance(out.realized, Device[out.device].buffer), f"device mismatch on realized got {type(out.realized)} expected {out.device}"
|
||||
|
||||
Reference in New Issue
Block a user