commaai
/
tinygrad
mirror of https://github.com/commaai/tinygrad.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

more ast_const for hardcoding consts [run_process_replay] (#6418)

This commit is contained in:
qazal 2024-09-09 11:35:08 +08:00 committed by GitHub
parent 25af78c593
commit 442150a8df
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 86 additions and 82 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
test/helpers.py
View File

@ -62,9 +62,10 @@ def assert_equiv_uops(u1:UOp, u2:UOp) -> None:
print_diff(u1, u2)
raise AssertionError("uops aren't equal.")
def ast_const(dtype:DType, val:ConstType, shape:Tuple[sint, ...]=(), st:Optional[ShapeTracker]=None) -> UOp:
st = st if st is not None else ShapeTracker.from_shape(()).reshape((1,)*len(shape)).expand(shape)
return UOp(UOps.CONST, dtype, (st.to_uop(),), dtypes.as_const(val, dtype))
def ast_const(dtype:DType, val:ConstType, shape:Tuple[sint, ...]=(), st:Optional[ShapeTracker]=None, st_src:Optional[Tuple[UOp]]=None) -> UOp:
if st_src is None:
st_src = (st.to_uop() if st is not None else ShapeTracker.from_shape(()).reshape((1,)*len(shape)).expand(shape).to_uop(),)
return UOp(UOps.CONST, dtype, st_src, dtypes.as_const(val, dtype))
T = TypeVar("T")
def timeit(fxn:Callable[..., T], *args, **kwargs) -> Tuple[T, float]:

25
test/test_linearizer_dumb.py
View File

@ -3,6 +3,7 @@
# like test_linearizer_failures, but they don't have to fail
import unittest
from test.helpers import ast_const
from tinygrad import Device, dtypes
from tinygrad.dtype import PtrDType
from tinygrad.ops import UOp, UOps, BinaryOps, TernaryOps
@ -30,9 +31,9 @@ class TestLinearizerDumb(unittest.TestCase):
UOp(UOps.LOAD, dtypes.half, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.half), arg=2, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(64, 1, 512, 7, 7, 512, 3, 3), strides=(0, 0, 4608, 0, 0, 9, 3, 1), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),)),)),
UOp(UOps.CONST, dtypes.half, arg=0.9999950000374996, src=(
ast_const(dtypes.half, 0.9999950000374996, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(64, 1, 512, 7, 7, 1, 1, 1), strides=(0, 0, 0, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
UOp(UOps.CONST, dtypes.half, arg=0.0, src=(
ast_const(dtypes.half, 0.0, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(64, 1, 512, 7, 7, 1, 1, 1), strides=(0, 0, 0, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),))
opts = [Opt(op=OptOps.TC, axis=2, amt=2), Opt(op=OptOps.UPCAST, axis=2, amt=0), Opt(op=OptOps.UNROLL, axis=1, amt=0)]
k = Kernel(ast, opts=Device["METAL"].renderer)
@ -62,13 +63,13 @@ class TestLinearizerDumb(unittest.TestCase):
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=2, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1000, 1), strides=(0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
UOp(UOps.CONST, dtypes.bool, arg=True, src=(
ast_const(dtypes.bool, True, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1000, 1), strides=(0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),
UOp(UOps.ALU, dtypes.int, arg=BinaryOps.ADD, src=(
UOp(UOps.REDUCE_AXIS, dtypes.int, arg=(BinaryOps.ADD, (1,)), src=(
UOp(UOps.CONST, dtypes.int, arg=-1, src=(
ast_const(dtypes.int, -1, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1001, 1999), strides=(0, 0), offset=0, mask=((0, 1001), (999, 1999)), contiguous=False), View(shape=(1000, 1000), strides=(1, 2000), offset=0, mask=None, contiguous=False))), src=()),)),)),
UOp(UOps.CONST, dtypes.int, arg=1000, src=(
ast_const(dtypes.int, 1000, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1000, 1), strides=(0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),)),))
opts = [Opt(op=OptOps.UNROLL, axis=0, amt=4), Opt(op=OptOps.LOCAL, axis=0, amt=8)]
k = Kernel(ast, opts=Device[Device.DEFAULT].renderer)
@ -113,14 +114,14 @@ class TestLinearizerDumb(unittest.TestCase):
UOp(UOps.ALU, dtypes.bool, arg=BinaryOps.CMPNE, src=(
UOp(UOps.ALU, dtypes.int, arg=BinaryOps.ADD, src=(
UOp(UOps.REDUCE_AXIS, dtypes.int, arg=(BinaryOps.ADD, (2,)), src=(
UOp(UOps.CONST, dtypes.int, arg=1, src=(
ast_const(dtypes.int, 1, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(32001, 63999), strides=(0, 0), offset=0, mask=((0, 32001), (31999, 63999)), contiguous=False), View(shape=(4096, 32000, 32000), strides=(0, 1, 64000), offset=0, mask=None, contiguous=False))), src=()),)),)),
UOp(UOps.CONST, dtypes.int, arg=-1, src=(
ast_const(dtypes.int, -1, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(4096, 32000, 1), strides=(0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
UOp(UOps.LOAD, dtypes.int, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.int), arg=1, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(4096, 32000, 1), strides=(0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
UOp(UOps.CONST, dtypes.bool, arg=True, src=(
ast_const(dtypes.bool, True, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(4096, 32000, 1), strides=(0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),
UOp(UOps.LOAD, dtypes.half, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.half), arg=2, src=()),
@ -148,7 +149,7 @@ class TestLinearizerDumb(unittest.TestCase):
UOp(UOps.LOAD, dtypes.int, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.int), arg=2, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(3, 1, 5), strides=(0, 0, 1), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),
UOp(UOps.CONST, dtypes.bool, arg=True, src=(
ast_const(dtypes.bool, True, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(3, 1, 5), strides=(0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=3, src=()),
@ -175,11 +176,11 @@ class TestLinearizerDumb(unittest.TestCase):
UOp(UOps.LOAD, dtypes.long, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.long), arg=1, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(3, 6), strides=(6, 1), offset=0, mask=None, contiguous=True),)), src=()),)),
UOp(UOps.CONST, dtypes.long, arg=-1, src=(
ast_const(dtypes.long, -1, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(3, 6), strides=(0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
UOp(UOps.CONST, dtypes.bool, arg=True, src=(
ast_const(dtypes.bool, True, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(3, 6), strides=(0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
UOp(UOps.CONST, dtypes.float, arg=0.0, src=(
ast_const(dtypes.float, 0.0, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(3, 6), strides=(0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=2, src=()),

116
test/test_linearizer_failures.py
View File

@ -8,7 +8,7 @@ from tinygrad import Device, dtypes, Tensor
from tinygrad.dtype import PtrDType
from tinygrad.helpers import CI
from test.external.fuzz_linearizer import compare_linearizer
from test.helpers import is_dtype_supported
from test.helpers import is_dtype_supported, ast_const
from tinygrad.shape.shapetracker import ShapeTracker
from tinygrad.shape.view import View
@ -93,9 +93,9 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.ALU, dtypes.float, arg=BinaryOps.ADD, src=(
x5:=UOp(UOps.ALU, dtypes.float, arg=BinaryOps.MUL, src=(
UOp(UOps.ALU, dtypes.float, arg=BinaryOps.ADD, src=(
UOp(UOps.CONST, dtypes.float, arg=0.1464405059814453, src=(
ast_const(dtypes.float, 0.1464405059814453, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(2, 1, 4, 1, 3, 1, 4, 1), strides=(0, 0, 0, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),
UOp(UOps.CONST, dtypes.float, arg=1.0, src=(
ast_const(dtypes.float, 1.0, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(2, 1, 4, 1, 3, 1, 4, 1), strides=(0, 0, 0, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=1, src=()),
@ -112,9 +112,9 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(10, 1), strides=(1, 0), offset=0, mask=None, contiguous=True),)), src=()),
UOp(UOps.ALU, dtypes.int, arg=BinaryOps.ADD, src=(
UOp(UOps.REDUCE_AXIS, dtypes.int, arg=(BinaryOps.ADD, (1,)), src=(
UOp(UOps.CONST, dtypes.int, arg=-1, src=(
ast_const(dtypes.int, -1, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(11, 19), strides=(0, 0), offset=0, mask=((0, 11), (9, 19)), contiguous=False), View(shape=(10, 10), strides=(1, 20), offset=0, mask=None, contiguous=False))), src=()),)),)),
UOp(UOps.CONST, dtypes.int, arg=10, src=(
ast_const(dtypes.int, 10, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(10, 1), strides=(0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),))
opts = [Opt(op=OptOps.UPCAST, axis=0, amt=2), Opt(op=OptOps.UPCAST, axis=0, amt=0)]
# COMPILE FAILED, KeyError: UOps.CONST
@ -152,9 +152,9 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=2, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1, 1, 4096), strides=(0, 0, 1), offset=0, mask=None, contiguous=True),)), src=()),)),)),
x9,)),)),
UOp(UOps.CONST, dtypes.float, arg=0.000244140625, src=(
ast_const(dtypes.float, 0.000244140625, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1, 1, 1), strides=(0, 0, 0), offset=0, mask=None, contiguous=True),)), src=()),)),)),
UOp(UOps.CONST, dtypes.float, arg=1e-06, src=(
ast_const(dtypes.float, 1e-06, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1, 1, 1), strides=(0, 0, 0), offset=0, mask=None, contiguous=True),)), src=()),)),)),)),)),)),))
opts = [Opt(op=OptOps.UNROLL, axis=0, amt=4), Opt(op=OptOps.UNROLL, axis=0, amt=4), Opt(op=OptOps.UNROLL, axis=0, amt=4), Opt(op=OptOps.UNROLL, axis=0, amt=4)]
# fatal error: bracket nesting level exceeded maximum of 256
@ -214,17 +214,17 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=2, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(512, 64, 6, 6), strides=(0, 1, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
UOp(UOps.CONST, dtypes.float, arg=0.0, src=(
ast_const(dtypes.float, 0.0, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(512, 64, 6, 6), strides=(0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=3, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(512, 64, 6, 6), strides=(0, 1, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
UOp(UOps.CONST, dtypes.float, arg=1.0, src=(
ast_const(dtypes.float, 1.0, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(512, 64, 6, 6), strides=(0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
UOp(UOps.ALU, dtypes.float, arg=BinaryOps.MUL, src=(
UOp(UOps.ALU, dtypes.float, arg=BinaryOps.MUL, src=(
UOp(UOps.ALU, dtypes.float, arg=BinaryOps.ADD, src=(
UOp(UOps.CONST, dtypes.float, arg=1.0, src=(
ast_const(dtypes.float, 1.0, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(512, 64, 3, 3, 2, 2), strides=(0, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False), View(shape=(512, 64, 3, 2, 3, 2), strides=(2304, 36, 12, 2, 4, 1), offset=0, mask=None, contiguous=False), View(shape=(512, 64, 6, 6), strides=(2304, 36, 6, 1), offset=0, mask=None, contiguous=True))), src=()),)),
UOp(UOps.CAST, dtypes.float, arg=None, src=(
UOp(UOps.ALU, dtypes.bool, arg=BinaryOps.CMPLT, src=(
@ -240,12 +240,12 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=2, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(512, 64, 6, 6), strides=(0, 1, 0, 0), offset=0, mask=None, contiguous=False), View(shape=(512, 64, 3, 3, 2, 2), strides=(2304, 36, 12, 2, 6, 1), offset=0, mask=None, contiguous=False), View(shape=(512, 64, 3, 2, 3, 2), strides=(2304, 36, 12, 2, 4, 1), offset=0, mask=None, contiguous=False), View(shape=(512, 64, 6, 6), strides=(2304, 36, 6, 1), offset=0, mask=None, contiguous=True))), src=()),)),)),
x42:=UOp(UOps.CONST, dtypes.float, arg=0.0, src=(
x42:=ast_const(dtypes.float, 0.0, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(512, 64, 6, 6), strides=(0, 0, 0, 0), offset=0, mask=None, contiguous=False), View(shape=(512, 64, 3, 3, 2, 2), strides=(2304, 36, 12, 2, 6, 1), offset=0, mask=None, contiguous=False), View(shape=(512, 64, 3, 2, 3, 2), strides=(2304, 36, 12, 2, 4, 1), offset=0, mask=None, contiguous=False), View(shape=(512, 64, 6, 6), strides=(2304, 36, 6, 1), offset=0, mask=None, contiguous=True))), src=()),)),)),
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=3, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(512, 64, 6, 6), strides=(0, 1, 0, 0), offset=0, mask=None, contiguous=False), View(shape=(512, 64, 3, 3, 2, 2), strides=(2304, 36, 12, 2, 6, 1), offset=0, mask=None, contiguous=False), View(shape=(512, 64, 3, 2, 3, 2), strides=(2304, 36, 12, 2, 4, 1), offset=0, mask=None, contiguous=False), View(shape=(512, 64, 6, 6), strides=(2304, 36, 6, 1), offset=0, mask=None, contiguous=True))), src=()),)),)),
UOp(UOps.CONST, dtypes.float, arg=1.0, src=(
ast_const(dtypes.float, 1.0, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(512, 64, 6, 6), strides=(0, 0, 0, 0), offset=0, mask=None, contiguous=False), View(shape=(512, 64, 3, 3, 2, 2), strides=(2304, 36, 12, 2, 6, 1), offset=0, mask=None, contiguous=False), View(shape=(512, 64, 3, 2, 3, 2), strides=(2304, 36, 12, 2, 4, 1), offset=0, mask=None, contiguous=False), View(shape=(512, 64, 6, 6), strides=(2304, 36, 6, 1), offset=0, mask=None, contiguous=True))), src=()),)),)),
UOp(UOps.ALU, dtypes.float, arg=UnaryOps.SQRT, src=(
UOp(UOps.CAST, dtypes.float, arg=None, src=(
@ -255,9 +255,9 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=4, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(64,), strides=(1,), offset=0, mask=None, contiguous=True), View(shape=(512, 64, 6, 6), strides=(0, 1, 0, 0), offset=0, mask=None, contiguous=False), View(shape=(512, 64, 3, 3, 2, 2), strides=(2304, 36, 12, 2, 6, 1), offset=0, mask=None, contiguous=False), View(shape=(512, 64, 3, 2, 3, 2), strides=(2304, 36, 12, 2, 4, 1), offset=0, mask=None, contiguous=False), View(shape=(512, 64, 6, 6), strides=(2304, 36, 6, 1), offset=0, mask=None, contiguous=True))), src=()),)),
UOp(UOps.CONST, dtypes.float, arg=5.425347222222222e-05, src=(
ast_const(dtypes.float, 5.425347222222222e-05, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(64,), strides=(0,), offset=0, mask=None, contiguous=False), View(shape=(512, 64, 6, 6), strides=(0, 1, 0, 0), offset=0, mask=None, contiguous=False), View(shape=(512, 64, 3, 3, 2, 2), strides=(2304, 36, 12, 2, 6, 1), offset=0, mask=None, contiguous=False), View(shape=(512, 64, 3, 2, 3, 2), strides=(2304, 36, 12, 2, 4, 1), offset=0, mask=None, contiguous=False), View(shape=(512, 64, 6, 6), strides=(2304, 36, 6, 1), offset=0, mask=None, contiguous=True))), src=()),)),)),
UOp(UOps.CONST, dtypes.float, arg=1e-05, src=(
ast_const(dtypes.float, 1e-05, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(64,), strides=(0,), offset=0, mask=None, contiguous=False), View(shape=(512, 64, 6, 6), strides=(0, 1, 0, 0), offset=0, mask=None, contiguous=False), View(shape=(512, 64, 3, 3, 2, 2), strides=(2304, 36, 12, 2, 6, 1), offset=0, mask=None, contiguous=False), View(shape=(512, 64, 3, 2, 3, 2), strides=(2304, 36, 12, 2, 4, 1), offset=0, mask=None, contiguous=False), View(shape=(512, 64, 6, 6), strides=(2304, 36, 6, 1), offset=0, mask=None, contiguous=True))), src=()),)),)),)),)),)),)),
x42,)),
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
@ -284,7 +284,7 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=1, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(2, 1, 4, 1, 3, 4, 2, 6, 1, 3), strides=(0, 0, 0, 0, 0, 18, 0, 3, 0, 1), offset=0, mask=None, contiguous=False),)), src=()),)),
UOp(UOps.CONST, dtypes.float, arg=1.0, src=(
ast_const(dtypes.float, 1.0, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(2, 1, 4, 1, 3, 4, 2, 6, 1, 3), strides=(0, 0, 0, 0, 0, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=2, src=()),
@ -307,7 +307,7 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=1, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(2, 1, 4, 1, 3, 4, 2, 6, 1, 3), strides=(0, 0, 0, 0, 0, 18, 0, 3, 0, 1), offset=0, mask=None, contiguous=False),)), src=()),)),
UOp(UOps.CONST, dtypes.float, arg=1.0, src=(
ast_const(dtypes.float, 1.0, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(2, 1, 4, 1, 3, 4, 2, 6, 1, 3), strides=(0, 0, 0, 0, 0, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=2, src=()),
@ -351,7 +351,7 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=1, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(2, 1, 4, 1, 3, 4, 2, 6, 1, 3), strides=(0, 0, 0, 0, 0, 18, 0, 3, 0, 1), offset=0, mask=None, contiguous=False),)), src=()),)),
UOp(UOps.CONST, dtypes.float, arg=1.0, src=(
ast_const(dtypes.float, 1.0, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(2, 1, 4, 1, 3, 4, 2, 6, 1, 3), strides=(0, 0, 0, 0, 0, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=2, src=()),
@ -390,7 +390,7 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=5, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1, 1, 112, 14, 14, 1, 1, 1), strides=(0, 0, 1, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),
UOp(UOps.CONST, dtypes.float, arg=1e-05, src=(
ast_const(dtypes.float, 1e-05, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1, 1, 112, 14, 14, 1, 1, 1), strides=(0, 0, 0, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),)),)),
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=6, src=()),
@ -409,7 +409,7 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=1, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1, 13, 1024), strides=(0, 1024, 1), offset=0, mask=None, contiguous=True),)), src=()),)),)),
UOp(UOps.CONST, dtypes.float, arg=0.0009765625, src=(
ast_const(dtypes.float, 0.0009765625, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1, 13, 1), strides=(0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),))
opts = [Opt(op=OptOps.GROUP, axis=0, amt=4), Opt(op=OptOps.UNROLL, axis=0, amt=0), Opt(op=OptOps.UNROLL, axis=0, amt=4), Opt(op=OptOps.GROUP, axis=0, amt=8), Opt(op=OptOps.UNROLL, axis=0, amt=4), Opt(op=OptOps.UNROLL, axis=1, amt=4)]
# COMPILE_ERROR on METAL/GPU (probably HIP/CUDA too) in fuzz_linearizer ast 154: bracket nesting level exceeded maximum of 256
@ -483,7 +483,7 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=1, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(4, 4), strides=(0, 1), offset=0, mask=None, contiguous=False),)), src=()),)),
UOp(UOps.CONST, dtypes.float, arg=1.0, src=(
ast_const(dtypes.float, 1.0, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(4, 4), strides=(0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),))
opts = [Opt(op=OptOps.UPCAST, axis=1, amt=0), Opt(op=OptOps.UPCAST, axis=0, amt=0)]
helper_test_lin(Kernel(ast), opts, failed_platforms=[])
@ -493,7 +493,7 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.STORE, None, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=0, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(45, 65), strides=(65, 1), offset=0, mask=None, contiguous=True),)), src=()),
UOp(UOps.CONST, dtypes.float, arg=1.0, src=(
ast_const(dtypes.float, 1.0, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(45, 65), strides=(0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),))
opts = [Opt(op=OptOps.PADTO, axis=0, amt=32)]
helper_test_lin(Kernel(ast), opts, failed_platforms=[])
@ -506,7 +506,7 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=0, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1, 96, 1, 1), strides=(0, 1, 0, 0), offset=0, mask=None, contiguous=True),)), src=()),
UOp(UOps.ALU, dtypes.float, arg=BinaryOps.MUL, src=(
x4:=UOp(UOps.CONST, dtypes.float, arg=0.000244140625, src=(
x4:=ast_const(dtypes.float, 0.000244140625, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1, 96, 1, 1), strides=(0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),
UOp(UOps.ALU, dtypes.float, arg=BinaryOps.MUL, src=(
UOp(UOps.ALU, dtypes.float, arg=BinaryOps.MUL, src=(
@ -580,7 +580,7 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=17, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1, 96, 1, 1), strides=(0, 1, 0, 0), offset=0, mask=None, contiguous=True),)), src=()),)),
UOp(UOps.CONST, dtypes.float, arg=2.0, src=(
ast_const(dtypes.float, 2.0, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1, 96, 1, 1), strides=(0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),)),
x80:=UOp(UOps.ALU, dtypes.float, arg=UnaryOps.RECIP, src=(
UOp(UOps.ALU, dtypes.float, arg=BinaryOps.ADD, src=(
@ -589,7 +589,7 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=18, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1, 96, 1, 1), strides=(0, 1, 0, 0), offset=0, mask=None, contiguous=True),)), src=()),)),
x4,)),
UOp(UOps.CONST, dtypes.float, arg=1e-05, src=(
ast_const(dtypes.float, 1e-05, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1, 96, 1, 1), strides=(0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),)),
x80,)),)),)),))
opts = []
@ -625,9 +625,9 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1024, 1), strides=(1, 0), offset=0, mask=None, contiguous=True),)), src=()),
UOp(UOps.ALU, dtypes.int, arg=BinaryOps.ADD, src=(
UOp(UOps.REDUCE_AXIS, dtypes.int, arg=(BinaryOps.ADD, (1,)), src=(
UOp(UOps.CONST, dtypes.int, arg=1, src=(
ast_const(dtypes.int, 1, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1025, 2047), strides=(0, 0), offset=0, mask=((0, 1025), (1023, 2047)), contiguous=False), View(shape=(1024, 1024), strides=(1, 2048), offset=0, mask=None, contiguous=False))), src=()),)),)),
UOp(UOps.CONST, dtypes.int, arg=-1, src=(
ast_const(dtypes.int, -1, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1024, 1), strides=(0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),))
opts = [Opt(op=OptOps.GROUP, axis=0, amt=16), Opt(op=OptOps.UNROLL, axis=0, amt=4)]
helper_test_lin(Kernel(ast), opts, failed_platforms=[])
@ -640,9 +640,9 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(128, 1), strides=(1, 0), offset=0, mask=None, contiguous=True),)), src=()),
UOp(UOps.ALU, dtypes.int, arg=BinaryOps.ADD, src=(
UOp(UOps.REDUCE_AXIS, dtypes.int, arg=(BinaryOps.ADD, (1,)), src=(
UOp(UOps.CONST, dtypes.int, arg=1, src=(
ast_const(dtypes.int, 1, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(129, 255), strides=(0, 0), offset=0, mask=((0, 129), (127, 255)), contiguous=False), View(shape=(128, 128), strides=(1, 256), offset=0, mask=None, contiguous=False))), src=()),)),)),
UOp(UOps.CONST, dtypes.int, arg=-1, src=(
ast_const(dtypes.int, -1, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(128, 1), strides=(0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),))
all_failing_opts = [
[Opt(op=OptOps.UPCAST, axis=0, amt=4), Opt(op=OptOps.GROUPTOP, axis=0, amt=32), Opt(op=OptOps.UNROLL, axis=0, amt=0)],
@ -699,17 +699,17 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.LOAD, dtypes.int, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.int), arg=1, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1,), strides=(0,), offset=0, mask=None, contiguous=True),)), src=()),)),)),
x9:=UOp(UOps.CONST, dtypes.bfloat16, arg=230.0, src=(
x9:=ast_const(dtypes.bfloat16, 230.0, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1,), strides=(0,), offset=0, mask=None, contiguous=True),)), src=()),)),)),
UOp(UOps.ALU, dtypes.bfloat16, arg=BinaryOps.ADD, src=(
UOp(UOps.ALU, dtypes.bfloat16, arg=BinaryOps.MUL, src=(
UOp(UOps.ALU, dtypes.bfloat16, arg=BinaryOps.MUL, src=(
x5,
UOp(UOps.CONST, dtypes.bfloat16, arg=0.004347826086956522, src=(
ast_const(dtypes.bfloat16, 0.004347826086956522, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1,), strides=(0,), offset=0, mask=None, contiguous=True),)), src=()),)),)),
UOp(UOps.CONST, dtypes.bfloat16, arg=0.199374800625, src=(
ast_const(dtypes.bfloat16, 0.199374800625, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1,), strides=(0,), offset=0, mask=None, contiguous=True),)), src=()),)),)),
UOp(UOps.CONST, dtypes.bfloat16, arg=1.99375e-07, src=(
ast_const(dtypes.bfloat16, 1.99375e-07, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1,), strides=(0,), offset=0, mask=None, contiguous=True),)), src=()),)),)),
UOp(UOps.ALU, dtypes.bfloat16, arg=BinaryOps.ADD, src=(
UOp(UOps.ALU, dtypes.bfloat16, arg=BinaryOps.MUL, src=(
@ -717,11 +717,11 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.ALU, dtypes.bfloat16, arg=BinaryOps.ADD, src=(
x5,
x9,)),
UOp(UOps.CONST, dtypes.bfloat16, arg=0.0012987012987012987, src=(
ast_const(dtypes.bfloat16, 0.0012987012987012987, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1,), strides=(0,), offset=0, mask=None, contiguous=True),)), src=()),)),)),
UOp(UOps.CONST, dtypes.bfloat16, arg=-0.19439062499999998, src=(
ast_const(dtypes.bfloat16, -0.19439062499999998, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1,), strides=(0,), offset=0, mask=None, contiguous=True),)), src=()),)),)),
UOp(UOps.CONST, dtypes.bfloat16, arg=0.199375, src=(
ast_const(dtypes.bfloat16, 0.199375, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1,), strides=(0,), offset=0, mask=None, contiguous=True),)), src=()),)),)),)),)),))
helper_test_lin(Kernel(ast), opts=[], failed_platforms=[])
@ -777,7 +777,7 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=2, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1, 16, 13, 13), strides=(0, 13, 1, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
UOp(UOps.CONST, dtypes.float, arg=1.4426950408889634, src=(
ast_const(dtypes.float, 1.4426950408889634, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1, 16, 13, 13), strides=(0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),)),)),))
opts = [Opt(op=OptOps.UNROLL, axis=0, amt=0), Opt(op=OptOps.PADTO, axis=1, amt=32)]
helper_test_lin(Kernel(ast), opts=opts, failed_platforms=[])
@ -817,28 +817,28 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.ALU, dtypes.float, arg=TernaryOps.WHERE, src=(
UOp(UOps.ALU, dtypes.bool, arg=BinaryOps.CMPNE, src=(
x5,
x10:=UOp(UOps.CONST, dtypes.float, arg=0.0, src=(
x10:=ast_const(dtypes.float, 0.0, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(32640,), strides=(0,), offset=0, mask=None, contiguous=False),)), src=()),)),)),
UOp(UOps.ALU, dtypes.float, arg=TernaryOps.WHERE, src=(
UOp(UOps.ALU, dtypes.bool, arg=BinaryOps.CMPLT, src=(
UOp(UOps.ALU, dtypes.float, arg=BinaryOps.ADD, src=(
UOp(UOps.ALU, dtypes.float, arg=BinaryOps.ADD, src=(
UOp(UOps.ALU, dtypes.float, arg=BinaryOps.MUL, src=(
UOp(UOps.CONST, dtypes.float, arg=0.06788442333021306, src=(
ast_const(dtypes.float, 0.06788442333021306, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(32640,), strides=(0,), offset=0, mask=((0, 26040),), contiguous=False),)), src=()),)),
x5,)),
UOp(UOps.CONST, dtypes.float, arg=-0.03394221166510653, src=(
ast_const(dtypes.float, -0.03394221166510653, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(32640,), strides=(0,), offset=0, mask=((0, 26040),), contiguous=False),)), src=()),)),)),
UOp(UOps.ALU, dtypes.float, arg=BinaryOps.ADD, src=(
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=2, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(32640,), strides=(1,), offset=-26040, mask=((26040, 32640),), contiguous=False),)), src=()),)),
UOp(UOps.CONST, dtypes.float, arg=-0.18257418583505536, src=(
ast_const(dtypes.float, -0.18257418583505536, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(32640,), strides=(0,), offset=0, mask=((26040, 32640),), contiguous=False),)), src=()),)),)),)),
x10,)),
UOp(UOps.CONST, dtypes.float, arg=-1.0, src=(
ast_const(dtypes.float, -1.0, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(32640,), strides=(0,), offset=0, mask=None, contiguous=False),)), src=()),)),
UOp(UOps.CONST, dtypes.float, arg=1.0, src=(
ast_const(dtypes.float, 1.0, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(32640,), strides=(0,), offset=0, mask=None, contiguous=False),)), src=()),)),)),
x10,)),)),)),)),))
opts = [Opt(op=OptOps.GROUPTOP, axis=0, amt=16)]
@ -859,7 +859,7 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=2, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(4, 1, 6, 10, 3, 1, 2, 5), strides=(0, 0, 10, 0, 0, 0, 5, 1), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),
UOp(UOps.CONST, dtypes.float, arg=0.0, src=(
ast_const(dtypes.float, 0.0, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(4, 1, 6, 10, 3, 1, 1, 1), strides=(0, 0, 0, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),))
opts = [Opt(op=OptOps.TC, axis=0, amt=2), Opt(op=OptOps.UNROLL, axis=0, amt=0)] if unroll else [Opt(op=OptOps.TC, axis=0, amt=2)]
helper_test_lin(Kernel(ast), opts=opts, failed_platforms=[])
@ -877,9 +877,9 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.ALU, dtypes.uint, arg=BinaryOps.ADD, src=(
UOp(UOps.REDUCE_AXIS, dtypes.uint, arg=(BinaryOps.ADD, (1,)), src=(
UOp(UOps.CAST, dtypes.uint, arg=None, src=(
UOp(UOps.CONST, dtypes.uchar, arg=1, src=(
ast_const(dtypes.uchar, 1, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(6, 9), strides=(0, 0), offset=0, mask=((0, 6), (4, 9)), contiguous=False), View(shape=(5, 5), strides=(1, 10), offset=0, mask=None, contiguous=False))), src=()),)),)),)),
UOp(UOps.CONST, dtypes.uint, arg=-1, src=(
ast_const(dtypes.uint, -1, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(5, 1), strides=(0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),)),))
opts = [Opt(op=OptOps.UPCAST, axis=0, amt=0)]
helper_test_lin(Kernel(ast), opts=opts, failed_platforms=[])
@ -907,7 +907,7 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=3, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(512, 1, 32, 24, 24, 1, 1, 1), strides=(0, 0, 1, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
UOp(UOps.CONST, dtypes.float, arg=0.0, src=(
ast_const(dtypes.float, 0.0, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(512, 1, 32, 24, 24, 1, 1, 1), strides=(0, 0, 0, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),))
for axis in [0,1,2,3,4,5]:
opts = [Opt(op=OptOps.TC, axis=axis, amt=2)]
@ -955,7 +955,7 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=3, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(10000, 1, 32, 24, 24, 1, 1, 1), strides=(0, 0, 1, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
UOp(UOps.CONST, dtypes.float, arg=0.0, src=(
ast_const(dtypes.float, 0.0, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(10000, 1, 32, 24, 24, 1, 1, 1), strides=(0, 0, 0, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),))
for axis in [0,1,2,3,4,5]:
opts = [Opt(op=OptOps.TC, axis=axis, amt=2)]
@ -970,9 +970,9 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(60000, 1), strides=(1, 0), offset=0, mask=None, contiguous=True),)), src=()),
UOp(UOps.ALU, dtypes.int, arg=BinaryOps.ADD, src=(
UOp(UOps.REDUCE_AXIS, dtypes.int, arg=(BinaryOps.ADD, (1,)), src=(
UOp(UOps.CONST, dtypes.int, arg=1, src=(
ast_const(dtypes.int, 1, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(60001, 119999), strides=(0, 0), offset=0, mask=((0, 60001), (59999, 119999)), contiguous=False), View(shape=(60000, 60000), strides=(1, 120000), offset=0, mask=None, contiguous=False))), src=()),)),)),
UOp(UOps.CONST, dtypes.int, arg=-1, src=(
ast_const(dtypes.int, -1, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(60000, 1), strides=(0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),))
for amt in [16,32]:
opts = [Opt(op=OptOps.GROUPTOP, axis=0, amt=amt), Opt(op=OptOps.UNROLL, axis=0, amt=0)]
@ -1066,11 +1066,11 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(2, 3, 2, 3, 1), strides=(0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),
UOp(UOps.ALU, dtypes.int, arg=BinaryOps.ADD, src=(
UOp(UOps.REDUCE_AXIS, dtypes.int, arg=(BinaryOps.ADD, (4,)), src=(
UOp(UOps.CONST, dtypes.int, arg=1, src=(
ast_const(dtypes.int, 1, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(3, 3), strides=(0, 0), offset=0, mask=((0, 3), (1, 3)), contiguous=False), View(shape=(2, 3, 2, 3, 3), strides=(0, 0, 1, 0, 4), offset=0, mask=((0, 2), (0, 3), (0, 2), (0, 3), (0, 2)), contiguous=False))), src=()),)),)),
x19:=UOp(UOps.CONST, dtypes.int, arg=-1, src=(
x19:=ast_const(dtypes.int, -1, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(2, 3, 2, 3, 1), strides=(0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),
x21:=UOp(UOps.CONST, dtypes.bool, arg=True, src=(
x21:=ast_const(dtypes.bool, True, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(2, 3, 2, 3, 1), strides=(0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
UOp(UOps.ALU, dtypes.bool, arg=BinaryOps.CMPNE, src=(
UOp(UOps.ALU, dtypes.bool, arg=BinaryOps.CMPNE, src=(
@ -1079,7 +1079,7 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(2, 3, 2, 3, 1), strides=(3, 1, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),
UOp(UOps.ALU, dtypes.int, arg=BinaryOps.ADD, src=(
UOp(UOps.REDUCE_AXIS, dtypes.int, arg=(BinaryOps.ADD, (4,)), src=(
UOp(UOps.CONST, dtypes.int, arg=1, src=(
ast_const(dtypes.int, 1, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(4, 5), strides=(0, 0), offset=0, mask=((0, 4), (2, 5)), contiguous=False), View(shape=(2, 3, 2, 3, 3), strides=(0, 0, 0, 1, 6), offset=0, mask=None, contiguous=False))), src=()),)),)),
x19,)),)),
x21,)),)),)),)),)),)),))
@ -1105,7 +1105,7 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.LOAD, dtypes.int, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.int), arg=2, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(512, 10), strides=(1, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
UOp(UOps.CONST, dtypes.bool, arg=True, src=(
ast_const(dtypes.bool, True, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(512, 10), strides=(0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
UOp(UOps.LOAD, dtypes.bool, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.bool), arg=3, src=()),
@ -1128,9 +1128,9 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(60000, 1), strides=(1, 0), offset=0, mask=None, contiguous=True),)), src=()),
UOp(UOps.ALU, dtypes.int, arg=BinaryOps.ADD, src=(
UOp(UOps.REDUCE_AXIS, dtypes.int, arg=(BinaryOps.ADD, (1,)), src=(
UOp(UOps.CONST, dtypes.int, arg=1, src=(
ast_const(dtypes.int, 1, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(60001, 119999), strides=(0, 0), offset=0, mask=((0, 60001), (59999, 119999)), contiguous=False), View(shape=(60000, 60000), strides=(1, 120000), offset=0, mask=None, contiguous=False))), src=()),)),)),
UOp(UOps.CONST, dtypes.int, arg=-1, src=(
ast_const(dtypes.int, -1, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(60000, 1), strides=(0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),))
opts = [Opt(op=OptOps.UPCAST, axis=0, amt=3)]
helper_test_lin(Kernel(ast), opts=opts, failed_platforms=[])
@ -1191,9 +1191,9 @@ class TestLinearizerFailures(unittest.TestCase):
UOp(UOps.LOAD, dtypes.int, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.int), arg=3, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1, 1, 20, 20, 20), strides=(0, 0, 0, 1, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
UOp(UOps.CONST, dtypes.bool, arg=True, src=(
ast_const(dtypes.bool, True, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1, 1, 20, 20, 20), strides=(0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),)),
UOp(UOps.CONST, dtypes.bool, arg=True, src=(
ast_const(dtypes.bool, True, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1, 1, 20, 1, 20), strides=(0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),))
opts = [Opt(op=OptOps.UPCAST, axis=1, amt=2)]
helper_test_lin(Kernel(ast, opts=Device[Device.DEFAULT].renderer), opts=opts, failed_platforms=[])

7
test/test_linearizer_overflows.py
View File

@ -1,5 +1,6 @@
# ruff: noqa: E501
import unittest
from test.helpers import ast_const
from tinygrad import dtypes, Device
from tinygrad.helpers import CI
from tinygrad.codegen.kernel import Kernel
@ -41,19 +42,19 @@ class TestLinearizerOverflow(unittest.TestCase):
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=2, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(64, 1, 64, 112, 112, 3, 7, 7), strides=(0, 0, 147, 0, 0, 49, 7, 1), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),
x16:=UOp(UOps.CONST, dtypes.float, arg=0.0, src=(
x16:=ast_const(dtypes.float, 0.0, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(64, 1, 64, 112, 112, 1, 1, 1), strides=(0, 0, 0, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=3, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(64, 1, 64, 112, 112, 1, 1, 1), strides=(0, 0, 1, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
UOp(UOps.ALU, dtypes.float, arg=UnaryOps.SQRT, src=(
UOp(UOps.ALU, dtypes.float, arg=BinaryOps.MUL, src=(
x23:=UOp(UOps.CONST, dtypes.float, arg=1.0, src=(
x23:=ast_const(dtypes.float, 1.0, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(64, 1, 64, 112, 112, 1, 1, 1), strides=(0, 0, 0, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),
UOp(UOps.ALU, dtypes.float, arg=UnaryOps.RECIP, src=(
UOp(UOps.ALU, dtypes.float, arg=BinaryOps.ADD, src=(
x23,
UOp(UOps.CONST, dtypes.float, arg=1e-05, src=(
ast_const(dtypes.float, 1e-05, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(64, 1, 64, 112, 112, 1, 1, 1), strides=(0, 0, 0, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),)),)),)),
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=4, src=()),

8
test/test_schedule.py
View File

@ -17,7 +17,7 @@ from tinygrad.helpers import AST_REWRITE, CI, DEBUG, FUSE_ARANGE, FUSE_CONV_BW,
from tinygrad.codegen.kernel import Kernel, verify_ast
from tinygrad.engine.schedule import create_schedule, get_output_st, reduceop_fusor, st_fixup, ScheduleItem
from tinygrad.engine.realize import CompiledRunner, run_schedule
from test.helpers import assert_equiv_uops, is_dtype_supported, Context, timeit
from test.helpers import assert_equiv_uops, ast_const, is_dtype_supported, Context, timeit
from tinygrad.lazy import LazyBuffer, view_supported_devices
from extra.models.llama import precompute_freqs_cis
@ -1666,7 +1666,7 @@ class TestScheduleRewrite(unittest.TestCase):
bufs = [UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.int), (), i) for i in range(2)]
ld = UOp(UOps.LOAD, dtypes.int, (bufs[1], ShapeTracker.from_shape((32, 32)).to_uop()))
r = UOp(UOps.REDUCE_AXIS, dtypes.int, (ld,), (BinaryOps.ADD, (0, 1)))
r = r + UOp(UOps.CONST, dtypes.int, (ShapeTracker.from_shape(()).to_uop(),), 2)
r = r + ast_const(dtypes.int, 2, ())
sink = UOp(UOps.SINK, None, (UOp(UOps.STORE, None, (bufs[0], ShapeTracker.from_shape(()).to_uop(), r)),))
rsink = graph_rewrite(sink, reduceop_fusor)
with self.assertRaisesRegex(AssertionError, "implicit reshape"): verify_ast(sink)
@ -1685,7 +1685,7 @@ class TestScheduleRewrite(unittest.TestCase):
bufs = [UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.int), (), i) for i in range(2)]
ld = UOp(UOps.LOAD, dtypes.int, (bufs[1], ShapeTracker.from_shape((32, 32)).to_uop()))
r = UOp(UOps.REDUCE_AXIS, dtypes.int, (ld,), (BinaryOps.ADD, (0, 1)))
for _ in range(24): r = r + UOp(UOps.CONST, dtypes.int, (ShapeTracker.from_shape(()).to_uop(),), 2)
for _ in range(24): r = r + ast_const(dtypes.int, 2, ())
sink = UOp(UOps.SINK, None, (UOp(UOps.STORE, None, (bufs[0], ShapeTracker.from_shape(()).to_uop(), r)),))
rsink, et = timeit(graph_rewrite, sink, reduceop_fusor)
with self.assertRaisesRegex(AssertionError, "implicit reshape"): verify_ast(sink)
@ -1701,7 +1701,7 @@ class TestScheduleRewrite(unittest.TestCase):
bufs = [UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.int), (), i) for i in range(2)]
ld = UOp(UOps.LOAD, dtypes.int, (bufs[1], ShapeTracker.from_shape((32, 32)).to_uop()))
r = UOp(UOps.REDUCE_AXIS, dtypes.int, (ld,), (BinaryOps.ADD, (0, 1)))
for _ in range(sz): r = r + UOp(UOps.CONST, dtypes.int, (ShapeTracker.from_shape(()).to_uop(),), 2)
for _ in range(sz): r = r + ast_const(dtypes.int, 2, ())
sink = UOp(UOps.SINK, None, (UOp(UOps.STORE, None, (bufs[0], ShapeTracker.from_shape(()).to_uop(), r)),))
rsink, et = timeit(graph_rewrite, sink, reduceop_fusor)
with self.assertRaisesRegex(AssertionError, "implicit reshape"): verify_ast(sink)

5
test/test_search.py
View File

@ -1,5 +1,6 @@
import unittest
from test.helpers import ast_const
from tinygrad.codegen.kernel import Opt, OptOps
from tinygrad.codegen.kernel import Kernel
from tinygrad.ops import UOp, UOps, BinaryOps
@ -49,7 +50,7 @@ class TestTimeLinearizer(unittest.TestCase):
UOp(UOps.STORE, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=0, src=()),
UOp(UOps.SHAPETRACKER, arg=ShapeTracker(views=(View(shape=(16,), strides=(1,), offset=0, mask=None, contiguous=True),))),
UOp(UOps.CONST, dtypes.float, arg=0.0, src=(
ast_const(dtypes.float, 0.0, st_src=(
UOp(UOps.SHAPETRACKER, arg=ShapeTracker(views=(View(shape=(16,), strides=(0,), offset=0, mask=None, contiguous=False),))),)),)),))
lin = Kernel(ast)
bufs = bufs_from_lin(lin)
@ -131,7 +132,7 @@ class TestBEAM(unittest.TestCase):
UOp(UOps.LOAD, dtypes.float, arg=None, src=(
UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=6, src=()),
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(384768,), strides=(1,), offset=-320640, mask=((320640, 384768),), contiguous=False), View(shape=(1, 501, 256), strides=(0, 1, 501), offset=256512, mask=None, contiguous=False))), src=()),)),)), # noqa: E501
UOp(UOps.CONST, dtypes.float, arg=1.4285714285714286, src=(
ast_const(dtypes.float, 1.4285714285714286, st_src=(
UOp(UOps.SHAPETRACKER, None, arg=ShapeTracker(views=(View(shape=(1, 501, 256), strides=(0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),)),)) # noqa: E501
lin = Kernel(ast)