All LazyOps in the Linearizer (#1905)

* loadop buffer on cpu

* works for GPU

* sort of working

* has bugs

* gpu tests pass

* fix some tests

* fix tensor cores

* fix test linearizer

* fix symbolic

* fix has_variable_shape

* non symbolic size

* disable weird test

* simple cache fix

* fix custom function

* fix kopt

* cleanups

* a bit broken on the assign

* contig check

* only buffer

* need that order

* idx
This commit is contained in:
George Hotz
2023-09-24 11:50:00 +08:00
committed by GitHub
parent 0f373b8b47
commit a5820390db
15 changed files with 151 additions and 138 deletions

View File

@@ -263,14 +263,13 @@ class Linearizer:
uops: List[UOp]
from tinygrad.tensor import Tensor
from tinygrad.helpers import prod
result = Tensor(2).realize() + Tensor(3).realize()
result.lazydata.realized = Device[Device.DEFAULT].buffer(prod(result.shape), result.dtype)
# use the real Linearizer to linearize 2+3
from tinygrad.lazy import _replace_loadops
from tinygrad.codegen.linearizer import Linearizer
from tinygrad.codegen.kernel import LinearizerOptions
linearizer = Linearizer(result.lazydata.op, result.lazydata, LinearizerOptions())
op, _ = _replace_loadops(result.lazydata.op)
linearizer = Linearizer(op)
linearizer.linearize()
# print the uops
@@ -279,13 +278,11 @@ for uop in linearizer.uops: print(uop)
# output:
"""
0 UOps.DEFINE_GLOBAL : ptr.dtypes.float [] ('data0', dtypes.float)
1 UOps.LOOP : [] ([], 'global')
2 UOps.LOOP : [] ([], 'local')
3 UOps.CONST : dtypes.float [] 2.0
4 UOps.CONST : dtypes.float [] 3.0
5 UOps.ALU : dtypes.float [3, 4] BinaryOps.ADD
6 UOps.STORE : [5] MemOp(name='data0', idx=<0>, local=False, memory_dtype=dtypes.float, valid=<1>, invalid_value=0.0)
7 UOps.ENDLOOP : [] ([], 'global+local')
1 UOps.CONST : dtypes.float [] 2.0
2 UOps.CONST : dtypes.float [] 3.0
3 UOps.ALU : dtypes.float [1, 2] BinaryOps.ADD
4 UOps.CONST : dtypes.int [] 0
5 UOps.STORE : [0, 4, 3] None
"""
# %%