diff --git a/test/helpers.py b/test/helpers.py index 59125aca..b441359c 100644 --- a/test/helpers.py +++ b/test/helpers.py @@ -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]: diff --git a/test/test_linearizer_dumb.py b/test/test_linearizer_dumb.py index e309287b..a337219e 100644 --- a/test/test_linearizer_dumb.py +++ b/test/test_linearizer_dumb.py @@ -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=()), diff --git a/test/test_linearizer_failures.py b/test/test_linearizer_failures.py index 7f5b0099..3bddf9c6 100644 --- a/test/test_linearizer_failures.py +++ b/test/test_linearizer_failures.py @@ -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=[]) diff --git a/test/test_linearizer_overflows.py b/test/test_linearizer_overflows.py index e4ec98b5..b1370724 100644 --- a/test/test_linearizer_overflows.py +++ b/test/test_linearizer_overflows.py @@ -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=()), diff --git a/test/test_schedule.py b/test/test_schedule.py index 8c448368..ccebfdca 100644 --- a/test/test_schedule.py +++ b/test/test_schedule.py @@ -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) diff --git a/test/test_search.py b/test/test_search.py index 2686a950..8f50a9f2 100644 --- a/test/test_search.py +++ b/test/test_search.py @@ -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)