From b6835f413421f360dd8db927d8a371e2e5b79b22 Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Wed, 8 Oct 2025 14:47:02 +0300 Subject: [PATCH] remove Ops.VIEW and related UOp methods (#12522) * remove Ops.VIEW and related UOp methods * update abstractions2.py * no ShapeTrackers in abstractions2.py * it's a size 1 --- docs/abstractions2.py | 8 +- test/opt/test_gen_float4.py | 1 + test/test_linearizer.py | 55 ------------ test/test_linearizer_dumb.py | 2 + test/test_schedule.py | 6 +- test/test_uops.py | 11 --- test/unit/test_tensor_uop_representation.py | 1 - test/unit/test_uop_spec.py | 97 --------------------- tinygrad/codegen/lowerer.py | 44 +--------- tinygrad/codegen/quantize.py | 14 +-- tinygrad/uop/__init__.py | 8 +- tinygrad/uop/ops.py | 39 ++------- tinygrad/uop/spec.py | 34 +------- tinygrad/viz/serve.py | 9 +- 14 files changed, 31 insertions(+), 298 deletions(-) delete mode 100644 test/unit/test_uop_spec.py diff --git a/docs/abstractions2.py b/docs/abstractions2.py index 1dc099c832..747b628644 100644 --- a/docs/abstractions2.py +++ b/docs/abstractions2.py @@ -42,7 +42,6 @@ import struct from tinygrad.dtype import dtypes from tinygrad.device import Buffer, Device from tinygrad.uop.ops import UOp, Ops -from tinygrad.shape.shapetracker import ShapeTracker # allocate some buffers + load in values out = Buffer(DEVICE, 1, dtypes.int32).allocate() @@ -51,13 +50,14 @@ b = Buffer(DEVICE, 1, dtypes.int32).allocate().copyin(memoryview(bytearray(struc # NOTE: a._buf is the same as the return from cpu.allocator.alloc # describe the computation +idx = UOp.const(dtypes.index, 0) buf_1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 1) buf_2 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 2) -ld_1 = UOp(Ops.LOAD, dtypes.int32, (buf_1.view(ShapeTracker.from_shape((1,))),)) -ld_2 = UOp(Ops.LOAD, dtypes.int32, (buf_2.view(ShapeTracker.from_shape((1,))),)) +ld_1 = UOp(Ops.LOAD, dtypes.int32, (buf_1.index(idx),)) +ld_2 = UOp(Ops.LOAD, dtypes.int32, (buf_2.index(idx),)) alu = ld_1 + ld_2 output_buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 0) -st_0 = UOp(Ops.STORE, dtypes.void, (output_buf.view(ShapeTracker.from_shape((1,))), alu)) +st_0 = UOp(Ops.STORE, dtypes.void, (output_buf.index(idx), alu)) s = UOp(Ops.SINK, dtypes.void, (st_0,)) # convert the computation to a "linearized" format (print the format) diff --git a/test/opt/test_gen_float4.py b/test/opt/test_gen_float4.py index 1b72514bfd..0b675eb469 100644 --- a/test/opt/test_gen_float4.py +++ b/test/opt/test_gen_float4.py @@ -149,6 +149,7 @@ class TestFloat4(unittest.TestCase): assert TestFloat4.count_float4(uops) == (1, 1) + @unittest.skip("Ops.VIEW no longer exists") def test_half4_load_unrolled(self): # from llama 7B shard 4 gpus ast = UOp(Ops.SINK, dtypes.void, arg=None, src=( diff --git a/test/test_linearizer.py b/test/test_linearizer.py index c7d92bcf28..c2ae2c3990 100644 --- a/test/test_linearizer.py +++ b/test/test_linearizer.py @@ -6,8 +6,6 @@ from tinygrad.codegen.opt import Opt, OptOps from tinygrad.codegen.gpudims import get_grouped_dims from tinygrad.uop.ops import UOp, Ops, GroupOp from tinygrad.device import Device, Buffer, is_dtype_supported -from tinygrad.shape.shapetracker import ShapeTracker -from tinygrad.shape.view import View from tinygrad.tensor import Tensor, _to_np_dtype from tinygrad.engine.realize import run_schedule, lower_schedule, CompiledRunner, get_program from tinygrad.helpers import Context, flatten, dedup, TC_SELECT, TC_OPT, RANGEIFY @@ -38,24 +36,6 @@ class TestLinearizer(unittest.TestCase): np.testing.assert_equal(a.numpy(), ta) np.testing.assert_equal(b.numpy(), tb) - def test_multioutput(self): - dtype, st = dtypes.int, ShapeTracker.from_shape((8,)) - g0, g1, g2, g3 = [UOp(Ops.DEFINE_GLOBAL, dtype.ptr(), arg=i) for i in range(4)] - a = UOp(Ops.LOAD, dtype, src=(g2.view(st),)) - b = UOp(Ops.LOAD, dtype, src=(g3.view(st),)) - out0 = UOp(Ops.STORE, dtypes.void, src=(g0.view(st), a + b)) - out1 = UOp(Ops.STORE, dtypes.void, src=(g1.view(st), a * b)) - sink = UOp(Ops.SINK, src=(out0, out1)) - - a_t = Tensor.full(st.shape, 2).contiguous().realize() - b_t = Tensor.full(st.shape, 3).contiguous().realize() - helper_linearizer_ast(sink, [a_t, b_t], wanna_output=[a_t.numpy()+b_t.numpy(), a_t.numpy()*b_t.numpy()]) - uops = get_program(sink, opts=[]).uops - stores = [u for u in uops if u.op is Ops.STORE] - mutable_bufs = dedup(flatten([[x for x in u.src[0].toposort() if x.op is Ops.DEFINE_GLOBAL] for u in stores])) - assert len(mutable_bufs) == len(stores) == 2 - self.assertSetEqual(set([u.arg for u in mutable_bufs]), set([0,1])) - def _test_no_nested_ranges(self, lins, skip=None): for l in lins: range_in_acc = flatten([[x for x in u.src if x.op is Ops.RANGE] for u in l.uops if u.op is Ops.DEFINE_REG]) @@ -437,41 +417,6 @@ class TestLinearizer(unittest.TestCase): # the global store doesn't change assert stores[1].src[1].dtype == dtypes.float - @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals") - @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4") - def test_skip_unmatching_upcasts(self): - Tensor.manual_seed(0) - c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(9600), arg=0, src=()) - c1 = c0.view(ShapeTracker(views=(View(shape=(240, 40, 1, 1), strides=(40, 1, 0, 0), offset=0, mask=None, contiguous=True),))) - c2 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(9600), arg=1, src=()) - c3 = c2.view(ShapeTracker(views=(View(shape=(240, 40, 1, 1), strides=(1, 240, 0, 0), offset=0, mask=None, contiguous=False),))) - c4 = c3.load() - c5 = c1.store(c4) - ast = c5.sink() - opt = [Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.LOCAL, axis=0, arg=16), - Opt(op=OptOps.LOCAL, axis=1, arg=2), Opt(op=OptOps.UPCAST, axis=3, arg=2)] - helper_linearizer_ast(ast, [Tensor.randn(240*40).realize()], opts=[opt]) - out = [u for u in get_program(ast, opts=opt).uops if u.op is Ops.STORE][0] - assert out.src[1].op is Ops.VECTORIZE and out.src[1].dtype == dtypes.float.vec(4) - - @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals") - @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4") - def test_skip_unmatching_upcasts_with_gep(self): - Tensor.manual_seed(0) - c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(256), arg=0, src=()) - c1 = c0.view(ShapeTracker(views=(View(shape=(8, 32, 1, 1), strides=(32, 1, 0, 0), offset=0, mask=None, contiguous=True),))) - c2 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(256), arg=1, src=()) - c3 = c2.view(ShapeTracker(views=(View(shape=(8, 32, 1, 1), strides=(1, 8, 0, 0), offset=0, mask=None, contiguous=False),))) - c4 = c3.load() - c5 = c1.store(c4) - ast = c5.sink() - opt = [Opt(op=OptOps.LOCAL, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=2, arg=2), Opt(op=OptOps.LOCAL, axis=1, arg=8), - Opt(op=OptOps.UPCAST, axis=1, arg=0), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.LOCAL, axis=0, arg=8), - Opt(op=OptOps.UPCAST, axis=1, arg=0), Opt(op=OptOps.UPCAST, axis=0, arg=2)] - helper_linearizer_ast(ast, [Tensor.randn(8*32).realize()], opts=[opt]) - out = [u for u in get_program(ast).uops if u.op is Ops.STORE][0] - assert out.src[1].op is Ops.VECTORIZE and out.src[1].dtype.count != 1 - # *** helpers *** def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]: diff --git a/test/test_linearizer_dumb.py b/test/test_linearizer_dumb.py index fd9b55ee2f..91d73218d3 100644 --- a/test/test_linearizer_dumb.py +++ b/test/test_linearizer_dumb.py @@ -32,6 +32,7 @@ class TestLinearizerFailure(unittest.TestCase): class TestLinearizerDumb(unittest.TestCase): @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "need local") + @unittest.skip("Ops.VALID no longer exists") def test_max_simplify_and_cancel(self): c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1000), arg=0, src=()) c1 = c0.view(ShapeTracker(views=(View(shape=(1000, 1), strides=(1, 0), offset=0, mask=None, contiguous=True),))) @@ -54,6 +55,7 @@ class TestLinearizerDumb(unittest.TestCase): # this was a bug in embedding, someday we should fold this anyway @unittest.skipUnless(is_dtype_supported(dtypes.half), f"half dtype not supported on {Device.DEFAULT}") + @unittest.skip("UOp.view is no longer supported") def test_llama_embedding(self): c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(4096), arg=0, src=()) c1 = c0.view(ShapeTracker(views=(View(shape=(4096, 1, 1), strides=(1, 0, 0), offset=0, mask=None, contiguous=True),))) diff --git a/test/test_schedule.py b/test/test_schedule.py index 876df4c2d3..b582b6e089 100644 --- a/test/test_schedule.py +++ b/test/test_schedule.py @@ -1954,8 +1954,7 @@ class TestSchedule(unittest.TestCase): a = Tensor([1,2,3,4]).realize() for _ in range(24): a = a + a new_uop = a.reshape(4,1).realize().uop - self.assertEqual(new_uop.st, ShapeTracker.from_shape((4,)).reshape((4, 1))) - self.assertEqual(swizzle_cnt(new_uop), 0) + assert new_uop.base.op is Ops.BUFFER @unittest.skipIf(CI and Device.DEFAULT == "NV", "crashes on NV CI") def test_limit_bufs_with_var(self): @@ -1981,9 +1980,6 @@ class TestSchedule(unittest.TestCase): sched = z.schedule() self.assertEqual(len(sched), kcount+1) -def swizzle_cnt(u:UOp) -> int: - return len([x for x in u.toposort() if x.op is Ops.VIEW and len(x.src) != 0 and x.src[0].op not in {Ops.BUFFER, Ops.DEFINE_GLOBAL, Ops.ASSIGN}]) - class TestSwizzle(unittest.TestCase): def test_swizzle_simple(self): Tensor.manual_seed(0) diff --git a/test/test_uops.py b/test/test_uops.py index 14315b040d..bb24c377c0 100644 --- a/test/test_uops.py +++ b/test/test_uops.py @@ -1,8 +1,6 @@ from typing import Optional, Any import unittest, math import numpy as np -from tinygrad.shape.shapetracker import ShapeTracker -from tinygrad.shape.view import View # noqa F401 from tinygrad.tensor import Tensor, _to_np_dtype from tinygrad.helpers import CI, DEBUG, getenv, Timing from tinygrad.dtype import dtypes, DType, AddrSpace @@ -492,15 +490,6 @@ class TestUOpMethod(unittest.TestCase): self.assertIs(x.replace(arg=None).arg, None) with self.assertRaises(AssertionError): x.replace(field="a") - def test_device(self): - x = UOp(Ops.VIEW, dtypes.int, (UOp.new_buffer(Device.DEFAULT, 1, dtypes.int), UOp.const(dtypes.int, 1)), ShapeTracker.from_shape(())) - self.assertEqual(x.device, Device.DEFAULT) - # NOTE: CONST doesn't have device - buffer, const = x.src - self.assertEqual(buffer.device, Device.DEFAULT) - self.assertEqual(const._device, None) - with self.assertRaises(AssertionError): const.device - class TestUOpStr(unittest.TestCase): def test_uop_str(self): a = UOp(Ops.CONST, dtypes.float, (), 2.0) + UOp(Ops.CONST, dtypes.float, (), 3.0) diff --git a/test/unit/test_tensor_uop_representation.py b/test/unit/test_tensor_uop_representation.py index f93ae2437f..9d53ae37e4 100644 --- a/test/unit/test_tensor_uop_representation.py +++ b/test/unit/test_tensor_uop_representation.py @@ -6,7 +6,6 @@ from tinygrad.uop.ops import UPat, Ops, UOp realized_pattern = UPat(Ops.BUFFER) # after realization, base tensor uops become RESHAPE(BUFFER) buffer_view_pattern = UPat(Ops.RESHAPE, src=(UPat(Ops.BUFFER),)) -const_pattern = UPat(Ops.CONST, src=(UPat(Ops.VIEW, src=(UPat(Ops.DEVICE),),))) def is_pattern_uop(u:UOp, pat:UPat): assert pat.match(u, {}), f"{u}\nis not\n{pat}" def is_pattern(ten:Tensor, pat:UPat): is_pattern_uop(ten.uop, pat) diff --git a/test/unit/test_uop_spec.py b/test/unit/test_uop_spec.py deleted file mode 100644 index 97f6d9040f..0000000000 --- a/test/unit/test_uop_spec.py +++ /dev/null @@ -1,97 +0,0 @@ -from __future__ import annotations -import unittest - -from tinygrad import Tensor -from tinygrad.helpers import DEBUG, RANGEIFY -from tinygrad.uop.ops import UOp, Ops, print_uops -from tinygrad.uop.spec import type_verify, ast_spec, tensor_uop_spec -from tinygrad.shape.shapetracker import ShapeTracker -from tinygrad import dtypes -from tinygrad.shape.view import View -from tinygrad.engine.realize import get_program -from tinygrad.device import Device - -class InvalidASTException(Exception): pass -def helper_test_verify_ast(*stores:UOp): - sink = UOp(Ops.SINK, dtypes.void, stores) - if DEBUG >= 3: - for op in stores: print(op) - try: type_verify(list(sink.toposort()), ast_spec) - except RuntimeError as e: raise InvalidASTException(e.args) - program = get_program(sink, Device[Device.DEFAULT].renderer) - - if DEBUG >= 6: print_uops(program.uops) - if DEBUG >= 4: print(program.src) - -class TestUOpSpec(unittest.TestCase): - def test_tiny_add(self): - dtype = dtypes.int - buf_0 = UOp(Ops.DEFINE_GLOBAL, dtype.ptr(), (), 0) - buf_1 = UOp(Ops.DEFINE_GLOBAL, dtype.ptr(), (), 1) - buf_2 = UOp(Ops.DEFINE_GLOBAL, dtype.ptr(), (), 2) - a = UOp(Ops.LOAD, dtype, (buf_1.view(ShapeTracker.from_shape((32, 1))),)) - b = UOp(Ops.LOAD, dtype, (buf_2.view(ShapeTracker.from_shape((32, 1))),)) - store = UOp(Ops.STORE, dtypes.void, (buf_0.view(ShapeTracker.from_shape((32, 1))), a+b)) - helper_test_verify_ast(store) - - def test_no_implicit_broadcasting(self): - bufs = [UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), i) for i in range(2)] - a = UOp(Ops.LOAD, dtypes.float, (bufs[1].view(ShapeTracker.from_shape((4, 32))),)) - b = a + UOp(Ops.REDUCE_AXIS, dtypes.float, (a,), (Ops.MAX, (1,))) - st = UOp(Ops.STORE, dtypes.void, (bufs[0].view(ShapeTracker.from_shape((4, 32))), b)) - with self.assertRaises(InvalidASTException): helper_test_verify_ast(st) - - def test_shrink_ok(self): - bufs = [UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), i) for i in range(2)] - a = UOp(Ops.LOAD, dtypes.float, (bufs[1].view(ShapeTracker((View((32, 32), strides=(32, 1), offset=0, mask=None, contiguous=True),))),)) - b = UOp(Ops.LOAD, dtypes.float, (bufs[1].view(ShapeTracker((View((32, 32), strides=(0, 1), offset=0, mask=None, contiguous=False),))),)) - st = UOp.store(bufs[0].view(ShapeTracker.from_shape((32, 32))), a+b) - helper_test_verify_ast(st) - - def test_reduce_store(self): - bufs = [UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), i) for i in range(2)] - a = UOp(Ops.LOAD, dtypes.float, (bufs[1].view(ShapeTracker.from_shape((32, 1))),)) - r = UOp(Ops.REDUCE_AXIS, dtypes.float, (a,), (Ops.ADD, (0,))) - st = UOp.store(bufs[0].view(ShapeTracker.from_shape((32, 1))), r) - with self.assertRaises(InvalidASTException): helper_test_verify_ast(st) - - def test_reduce_add_store(self): - bufs = [UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), i) for i in range(2)] - a = UOp(Ops.LOAD, dtypes.float, (bufs[1].view(ShapeTracker.from_shape((32, 1))),)) - r = UOp(Ops.REDUCE_AXIS, dtypes.float, (a,), (Ops.ADD, (0,))) - st = UOp.store(bufs[0].view(ShapeTracker.from_shape((32, 1))), r+a) - with self.assertRaises(InvalidASTException): helper_test_verify_ast(st) - - def test_assert_swizzle(self): - buf = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0) - a = UOp(Ops.LOAD, dtypes.float, (buf.view(ShapeTracker.from_shape((32, 1))),)) - r = UOp(Ops.REDUCE_AXIS, dtypes.float, (a,), (Ops.ADD, (0,))) - st = UOp.store(buf.view(ShapeTracker.from_shape((32, 1))), r.view(r.st.expand((32, 1)))+a) - with self.assertRaisesRegex(InvalidASTException, "UOp verification failed"): helper_test_verify_ast(st) - - def test_const_view_always_valid(self): - buf = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0) - a = UOp.const(dtypes.int, 0).replace(src=(UOp(Ops.VIEW, dtypes.void, (), ShapeTracker.from_shape(())),)) - st = UOp.store(buf.view(ShapeTracker.from_shape(())), a.cast(dtypes.float)) - helper_test_verify_ast(st) - - @unittest.skipIf(RANGEIFY, "RANGEIFY does not push views") - def test_assert_masked_view_in_const(self): - t = Tensor(6).uop - a = t.replace(src=(t.src[0].replace(arg=t.st.reshape((1,)).pad(((0, 1),))),)) - with self.assertRaisesRegex(RuntimeError, "UOp verification failed"): - type_verify([a], tensor_uop_spec) - -class TestUOpSink(unittest.TestCase): - def test_0(self): - s = UOp.sink() - self.assertEqual(len(s.src), 0) - - def test_1(self): - a = UOp.const(dtypes.int, 0) - s1 = UOp.sink(a) - s2 = a.sink() - self.assertIs(s1, s2) - -if __name__ == '__main__': - unittest.main() diff --git a/tinygrad/codegen/lowerer.py b/tinygrad/codegen/lowerer.py index 236aff36a4..06794d88c9 100644 --- a/tinygrad/codegen/lowerer.py +++ b/tinygrad/codegen/lowerer.py @@ -1,6 +1,6 @@ # the job of the lowerer is to do indexing from dataclasses import dataclass -from tinygrad.uop.ops import KernelInfo, UOp, Ops, PatternMatcher, UPat, sint_to_uop, AxisType, graph_rewrite, resolve +from tinygrad.uop.ops import KernelInfo, UOp, Ops, PatternMatcher, UPat, sint_to_uop, AxisType, graph_rewrite # ***** indexing ***** @@ -15,8 +15,8 @@ def shape_to_idx(s, axis_types, start=0): def get_index(ast:UOp) -> IndexContext: axis_types = ast.arg.axis_types if isinstance(ast.arg, KernelInfo) else () - if len(ast.full_shape) != len(axis_types) and ast.st is not None: - axis_types = tuple([AxisType.REDUCE if resolve(s != fs) else AxisType.LOOP for s,fs in zip(ast.shape, ast.full_shape)]) + #if len(ast.full_shape) != len(axis_types) and ast.st is not None: + # axis_types = tuple([AxisType.REDUCE if resolve(s != fs) else AxisType.LOOP for s,fs in zip(ast.shape, ast.full_shape)]) return IndexContext(axis_types, [], 0) # ***** lowering (given index) ***** @@ -26,29 +26,6 @@ def subblock(ctx: IndexContext, full_new_idx: list[UOp], src: UOp): ctx.start = lc.start return graph_rewrite(src, pm_lowerer, lc, name="subblock", bottom_up=True) -def lower_reduce_axis(ctx: IndexContext, x: UOp): - new_idxs = shape_to_idx(x.src[0].shape, ctx.axis_types, ctx.start) - full_new_idx = list(ctx.idxs) - for a in x.axis_arg: full_new_idx[a] = new_idxs[a] - ret = subblock(ctx, full_new_idx, x.src[0]) - return UOp(Ops.REDUCE, x.dtype, (ret,)+tuple([full_new_idx[i] for i in x.axis_arg]), x.arg[0]) - -def lower_store(ctx: IndexContext, x: UOp, buf: UOp): - # TODO: reenable after REDUCE_AXIS is fixed - #assert x.src[1].shape == x.src[0].shape, f"shape mismatch on store {x.src[1].shape} != {x.src[0].shape}" - - new_idxs = shape_to_idx(x.src[0].shape, ctx.axis_types, ctx.start) - idx = x.st_arg.to_valid_uop(new_idxs) - used_idxs = [x for x in idx.toposort() if x in new_idxs] - real_new_idxs = [] - for i in range(len(x.src[0].shape)): - if new_idxs[i] in used_idxs or len(ctx.idxs) <= i: real_new_idxs.append(new_idxs[i]) - else: real_new_idxs.append(ctx.idxs[i]) - - stored = subblock(ctx, real_new_idxs, x.src[1]) - used_ranges = [x for x in used_idxs if x.op is Ops.RANGE] - return buf.index(idx).store(stored, *used_ranges) - def fixup_wmma(ctx:IndexContext, x:UOp): if x.tag is not None: return None new_idxs = shape_to_idx(x.src[0].shape, ctx.axis_types, ctx.start) @@ -63,21 +40,6 @@ def fixup_wmma(ctx:IndexContext, x:UOp): return x.replace(src=srcs, arg=x.arg[:-2]+(new_x_arg_m2, new_x_arg_m1), tag=1) pm_lowerer = PatternMatcher([ - # TODO: remove these hacks - # hack for old style CONST(VIEW) (now it's just VIEW(CONST)) - (UPat((Ops.DEFINE_VAR, Ops.CONST), src=(UPat(Ops.VIEW, name="v"),), name="c"), lambda c,v: c.replace(src=()).view(v.arg)), - # hack for old style VALID (now it's just VIEW(CONST)) - (UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar("c"), UPat(Ops.CONST, arg=0)), lambda c,v: c.replace(src=()).view(v.arg)), - - # consts and loads - (UPat(Ops.VIEW, src=(UPat((Ops.CONST, Ops.DEFINE_VAR), name="c"),), name="view"), - lambda ctx,view,c: c if all(x.mask is None for x in view.arg.views) else view.arg.to_valid_uop(ctx.idxs).get_valid().where(c, c.const_like(0))), - (UPat(Ops.LOAD, src=(UPat.var("buf").view(),), allow_any_len=True, name="x"), - lambda ctx,buf,x: UOp(Ops.LOAD, x.dtype, (buf.index(x.st_arg.to_valid_uop(ctx.idxs)),)+x.src[1:])), - - # reduce/view_const - (UPat(Ops.REDUCE_AXIS, name="x"), lower_reduce_axis), - (UPat(Ops.STORE, src=(UPat.var("buf").view(),), allow_any_len=True, name="x"), lower_store), (UPat(Ops.WMMA, name="x"), fixup_wmma), # axis fixups for WMMA diff --git a/tinygrad/codegen/quantize.py b/tinygrad/codegen/quantize.py index a94bec18bb..ef34462c22 100644 --- a/tinygrad/codegen/quantize.py +++ b/tinygrad/codegen/quantize.py @@ -27,13 +27,13 @@ pm_quant = symbolic+PatternMatcher([ (UPat.var("x")*UPat.cvar("c1", dtype=dtypes.floats) + UPat.var("y")*UPat.cvar("c2", dtype=dtypes.floats), lambda x,y,c1,c2: (x+y)*c1 if abs(c1.arg-c2.arg) < 1e-9 else None), # mul 0 * c1 is 0 - (UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar("c1"), UPat(Ops.CONST, arg=0)) * - UPat(Ops.LOAD, src=(UPat().view(name="v"),)).cast(dtypes.int).cast(dtypes.float).named("ld"), lambda ld,v,c1: ld*c1), + #(UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar("c1"), UPat(Ops.CONST, arg=0)) * + # UPat(Ops.LOAD, src=(UPat().view(name="v"),)).cast(dtypes.int).cast(dtypes.float).named("ld"), lambda ld,v,c1: ld*c1), # mul (with plus) 0 * c1 is 0 - (UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar("c1"), UPat(Ops.CONST, arg=0)) * - (UPat(Ops.LOAD, src=(UPat().view(name="v"),)).cast(dtypes.int) + \ - UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar(), UPat(Ops.CONST, arg=0))).cast(dtypes.float).named("ld"), - lambda ld,v,c1: ld*c1), + #(UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar("c1"), UPat(Ops.CONST, arg=0)) * + # (UPat(Ops.LOAD, src=(UPat().view(name="v"),)).cast(dtypes.int) + \ + # UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar(), UPat(Ops.CONST, arg=0))).cast(dtypes.float).named("ld"), + # lambda ld,v,c1: ld*c1), # const push through add ((UPat.var("x")*UPat.cvar("c1") + UPat.var("y")*UPat.cvar("c2")) * UPat.cvar("c3"), lambda x,y,c1,c2,c3: (x*c1*c3) + (y*c2*c3)), @@ -64,4 +64,4 @@ pm_quant = symbolic+PatternMatcher([ lambda v1,v2,c1,r: r.replace(src=(v1*v2,)) + r.replace(src=(c1*v2,))), (UPat(Ops.REDUCE_AXIS, src=((UPat(Ops.CAST, name="v1")+UPat.var("c1")) * (UPat(Ops.CAST, name="v2",)+UPat.var("c2")),), name="r"), lambda v1,v2,c1,c2,r: r.replace(src=(v1*v2,)) + r.replace(src=(c2*v1,)) + r.replace(src=(c1*v2,)) + r.replace(src=(c1*c2,))), -]) \ No newline at end of file +]) diff --git a/tinygrad/uop/__init__.py b/tinygrad/uop/__init__.py index dad8229d8f..480971ddb4 100644 --- a/tinygrad/uop/__init__.py +++ b/tinygrad/uop/__init__.py @@ -33,12 +33,6 @@ class Ops(FastEnum): RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); PAD = auto(); SHRINK = auto(); FLIP = auto() # noqa: E702 MULTI = auto() # MULTI is really a movement op - # view is what all movement ops become - VIEW = auto() - - # TODO: remove VALID with the VIEW(CONST(DEVICE)) refactor - VALID = auto() - # TODO: unify these ops into the levels of the memory hierarchy. depends on ASSIGN is STORE DEFINE_GLOBAL = auto(); DEFINE_LOCAL = auto(); DEFINE_REG = auto() # noqa: E702 @@ -100,7 +94,7 @@ class GroupOp: Irreducible = {Ops.CONST, Ops.DEFINE_VAR, Ops.SPECIAL, Ops.RANGE} Movement = {Ops.RESHAPE, Ops.EXPAND, Ops.PERMUTE, Ops.PAD, Ops.SHRINK, Ops.FLIP} - Buffer = {Ops.LOAD, Ops.STORE, Ops.VALID, Ops.CONST, Ops.DEFINE_VAR} + Buffer = {Ops.LOAD, Ops.STORE, Ops.CONST, Ops.DEFINE_VAR} Block = {Ops.BLOCK, Ops.BLOCKEND, Ops.BLOCKSTART} # BinaryOps that can be flipped diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 98b28cc9fa..379a780a82 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -7,7 +7,7 @@ from tinygrad.uop import Ops, GroupOp from tinygrad.uop.mathtraits import MathTrait from tinygrad.dtype import ConstType, ImageDType, dtypes, DType, truncate, PtrDType, least_upper_dtype, Invalid, InvalidType from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten, TRACEMETA -from tinygrad.helpers import PICKLE_BUFFERS, PROFILE, dedup, cdiv, cmod, diskcache_put, to_function_name, cpu_profile, TracingKey, RANGEIFY, VIZ, SPEC +from tinygrad.helpers import PICKLE_BUFFERS, PROFILE, dedup, cdiv, cmod, diskcache_put, to_function_name, cpu_profile, TracingKey, VIZ, SPEC from tinygrad.helpers import strip_parens if TYPE_CHECKING: from tinygrad.shape.shapetracker import ShapeTracker @@ -186,8 +186,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass): if self.op is Ops.BARRIER: return None if self.op in GroupOp.Block: return None from tinygrad.shape.shapetracker import ShapeTracker - # VIEW and MovementOps define a new ShapeTracker from the arg - if self.op is Ops.VIEW: return self.arg + # MovementOps define a new ShapeTracker from the arg if self.op is Ops.BUFFERIZE: return ShapeTracker.from_shape(tuple([int(r.vmax+1) for r in self.src[1:]])) # allow reshape from nothing if self.op is Ops.RESHAPE and self.src[0].st is None: return ShapeTracker.from_shape(self.arg) @@ -198,7 +197,6 @@ class UOp(MathTrait, metaclass=UOpMetaClass): if self.op is Ops.STORE and self.dtype is not dtypes.void: return self.src[0].src[0].st # BufferOps and ASSIGN flow ShapeTracker from a direct edge if self.op in {Ops.STORE, Ops.ASSIGN, Ops.LOAD}: return self.src[0].st - if self.op in GroupOp.Buffer: return views[0] if (views:=[x.st for x in self.src if x.op is Ops.VIEW]) else None # BUFFER/BUFFER_VIEW and KERNEL only have a size if self.op in {Ops.BUFFER, Ops.BUFFER_VIEW}: return ShapeTracker.from_shape((self.size,)) @@ -229,12 +227,6 @@ class UOp(MathTrait, metaclass=UOpMetaClass): case _: shape = src_sts[0].shape return ShapeTracker.from_shape(shape) - @functools.cached_property - def full_shape(self) -> tuple[sint, ...]: - if self.op is Ops.VIEW: return self.shape - # NOTE: if a parent doesn't have st its full_shape is empty - parent_shapes = [x.full_shape for x in self.src] - return tuple(smax(x) for x in itertools.zip_longest(*parent_shapes, fillvalue=1)) @property def shape(self) -> tuple[sint, ...]: assert self.st is not None, f"{self.op} doesn't have a shape" @@ -345,17 +337,8 @@ class UOp(MathTrait, metaclass=UOpMetaClass): if isinstance(b, UOp): return b.unbind()[0] if b.op is Ops.BIND else b if isinstance(b, tuple) and all_same(b): b = b[0] # doesn't have to be a VCONST if they are all the same ret = UOp(Ops.VCONST if isinstance(b, tuple) else Ops.CONST, dtype, arg=dtypes.as_const(b, dtype), src=() if src is None else (src,)) - if RANGEIFY: - # VIEW on const is no longer supported in RANGEIFY - if device is not None: ret = ret.replace(src=(UOp(Ops.DEVICE, arg=device),)) - if shape is not None: ret = ret.reshape((1,)*len(shape)).expand(shape) - else: - if shape is not None: - from tinygrad.shape.shapetracker import ShapeTracker - ret = ret.replace(src=(UOp(Ops.VIEW, dtypes.void, (), ShapeTracker.from_shape(shape, (0,)*len(shape))),)) - if device is not None: - if shape is not None: ret = ret.replace(src=(UOp(Ops.DEVICE, arg=device).view(unwrap(ret.st)),)) - else: ret = ret.replace(src=(UOp(Ops.DEVICE, arg=device),)) + if device is not None: ret = ret.replace(src=(UOp(Ops.DEVICE, arg=device),)) + if shape is not None: ret = ret.reshape((1,)*len(shape)).expand(shape) return ret @staticmethod def range(end:sint, *arg): @@ -456,10 +439,9 @@ class UOp(MathTrait, metaclass=UOpMetaClass): @property def base(self) -> UOp: - if (self.op is Ops.VIEW and len(self.src) != 0) or self.op in GroupOp.Movement: return self.src[0].base + if self.op in GroupOp.Movement: return self.src[0].base if self.op is Ops.MULTI: return self.src[0].base # MULTI is really a VIEW return self - def view(self, new_st:ShapeTracker) -> UOp: return UOp(Ops.VIEW, self.dtype, (self,), new_st) def _mop(self, op:Ops, arg) -> UOp: ret = UOp(op, self.dtype, (self,), arg) @@ -572,8 +554,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass): all_vars = set([x for x in self.toposort() if x.op is Ops.DEFINE_VAR]) return bound_vars.union(set([x for x in all_vars if x not in bound_var_base])) def variables(self) -> list[Variable]: - st_vars: list[set[Variable]] = [x.arg.vars() for x in self.toposort() if x.op is Ops.VIEW] - return sorted(set.union(*st_vars, set([x.unbind()[0] if x.op is not Ops.DEFINE_VAR else x for x in self.vars()])), key=lambda v: v.arg) + return sorted(set([x.unbind()[0] if x.op is not Ops.DEFINE_VAR else x for x in self.vars()]), key=lambda v: v.arg) # *** uop symbolic stuff *** @@ -790,7 +771,6 @@ class UPat(MathTrait): # copied from UOp def sink(self, *srcs:UPat|None, **kwargs): return UPat(Ops.SINK, dtypes.void, (self,)+tuple([x for x in srcs if x is not None]), **kwargs) def index(self, idx:UPat, valid:UPat|None=None): return UPat(Ops.INDEX, self.dtype, (self,idx,valid) if valid is not None else (self,idx)) - def view(self, st=None, **kwargs): return UPat(Ops.VIEW, self.dtype, (self,), st, **kwargs) def cast(self, dtype=None, **kwargs): return UPat(Ops.CAST, dtype, (self,), **kwargs) def bitcast(self, dtype=None): return UPat(Ops.BITCAST, dtype, (self,)) def gep(self, i:int|None=None, **kwargs): return UPat(Ops.GEP, None, (self,), (i,) if i is not None else None, **kwargs) @@ -1168,7 +1148,6 @@ renderer = PatternMatcher([ (UPat(Ops.MULACC, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"({x.src[0].arg}*{x.src[1].arg}+{x.src[2].arg})")), (UPat(Ops.WHERE, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"({x.src[1].arg} if {x.src[0].arg} else {x.src[2].arg})")), (UPat(set(syms.keys()), src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"({x.src[0].arg}{syms[x.op]}{x.src[1].arg})")), - (UPat(Ops.VIEW, src=(UPat(Ops.NOOP),), name="x"), lambda x: UOp(Ops.NOOP, arg=f"{x.src[0].arg}.view({x.arg})")), (UPat((Ops.INDEX, Ops.BUFFERIZE), name="x"), lambda x: UOp(Ops.NOOP, arg=''.join([f"[{strip_parens(y.arg)}]" for y in x.src[1:]])) if all(y.op is Ops.NOOP for y in x.src[1:]) else None), (UPat(Ops.VECTORIZE, src=UPat(Ops.NOOP), name="x"), @@ -1195,8 +1174,6 @@ pm_pyrender = PatternMatcher([ arg=f"{x.src[0].arg}.{sugar[x.op]}({', '.join([y.arg for y in x.src[1:]] + ([f'arg={str(x.arg)}'] if x.arg is not None else []))})")), (UPat(Ops.REDUCE_AXIS, src=(UPat(Ops.NOOP),), name="x"), lambda x: UOp(Ops.NOOP, arg=f"{x.src[0].arg}.f({x.op}, arg=({', '.join([str(y) for y in x.arg])}))")), - (UPat(Ops.VALID, src=(UPat(Ops.NOOP),), name="x"), - lambda x: UOp(Ops.NOOP, arg=f"{x.src[0].arg}.f({x.op}, dtype=dtypes.bool)")), ]) @Context(SPEC=0) @@ -1205,8 +1182,8 @@ def pyrender(ast:UOp) -> list[str]: to_render = set() for u in ast.toposort(): if u.op is Ops.STORE: to_render.add(u.src[1]) - if len(cmap[u]) == 1 and u.op not in {Ops.DEFINE_GLOBAL, Ops.VIEW, Ops.LOAD} or u.op in {Ops.CONST}: continue - if u.op in {Ops.SINK, Ops.VIEW}: + if len(cmap[u]) == 1 and u.op not in {Ops.DEFINE_GLOBAL, Ops.LOAD} or u.op in {Ops.CONST}: continue + if u.op in {Ops.SINK}: for s in u.src: to_render.add(s) to_render.add(u) ret: list[str] = [] diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 262d960771..9b7dd85497 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -1,8 +1,7 @@ from typing import cast, Callable from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, print_uops, python_alu, graph_rewrite, AxisType from tinygrad.dtype import DType, ImageDType, dtypes, PtrDType, AddrSpace, Invalid -from tinygrad.helpers import all_same, prod, DEBUG, IGNORE_OOB, Context, cpu_profile, RANGEIFY -from tinygrad.shape.shapetracker import ShapeTracker +from tinygrad.helpers import all_same, prod, DEBUG, IGNORE_OOB, Context, cpu_profile try: import z3 # older versions of z3 dont have some operators like & overloaded @@ -64,8 +63,6 @@ buffer_spec = PatternMatcher([ (UPat(Ops.BUFFER_VIEW, src=(UPat(Ops.BUFFER),), name="buf_view"), lambda buf_view: isinstance(buf_view.arg, tuple) and len(buf_view.arg) == 2 and all(isinstance(arg, (int, UOp)) for arg in buf_view.arg)), (UPat(Ops.BUFFER_VIEW, src=(UPat(Ops.MSTACK, src=UPat(Ops.BUFFER)),)), lambda: True), - # allow VIEW here. TODO: what views specifically are allowed? does this mess with gradient? - (UPat(Ops.VIEW), lambda: True), ]) assign_spec = PatternMatcher([ @@ -92,17 +89,10 @@ tensor_uop_spec = buffer_spec+assign_spec+PatternMatcher([ # this is fine as long as it's a realized buffer or const and base dtypes match. ((isinstance(mv.dtype, ImageDType) or isinstance(x.dtype, ImageDType)) and x.dtype.base == mv.dtype.base \ and x.base.op in {Ops.BUFFER,Ops.ASSIGN,Ops.CONST})), - (UPat(Ops.VIEW, src=(UPat.var("x"),)), lambda x: x.base.op in {Ops.BUFFER, Ops.BUFFER_VIEW, Ops.ASSIGN, Ops.CONST, Ops.DEVICE}), # Tensor variable bindings (UPat(Ops.BIND, (dtypes.int,dtypes.index,), (UPat(Ops.DEFINE_VAR), UPat.cvar(dtype=(dtypes.int,dtypes.index,))), arg=None), lambda: True), - # Tensor const has a device and an unmasked ShapeTracker of stride 0 - # NOTE: variables in shape can cause multiple views in this ShapeTracker and other issues, see TestSymbolicJit.test_ones_sum - # TODO: remove after rangeify is default - (UPat(Ops.CONST, src=(UPat.any(UPat(Ops.VIEW, src=(UPat(Ops.DEVICE),), name="st"), - UPat(Ops.VIEW, src=(UPat(Ops.DEVICE), UPat(Ops.BIND)), name="st")),)), - lambda st: len(st.st.views) == 1 and all(v.mask is None for v in st.st.views)), (UPat(Ops.CONST, src=(UPat(Ops.DEVICE),)), lambda: True), # DETACH and CONTIGUOUS change how we interpret the source UOp @@ -167,20 +157,8 @@ spec = PatternMatcher([ all(isinstance(ra, int) for ra in rng.arg[0:-1]) and isinstance(rng.arg[-1], AxisType)), (UPat(Ops.SPECIAL, src=(UPat.var("x"),), name="s"), lambda s,x: s.dtype == x.dtype == dtypes.int32 and isinstance(s.arg, str)), - (UPat(Ops.VIEW, dtypes.void, src=(), name="x"), lambda x: isinstance(x.arg, ShapeTracker)), - (UPat(Ops.VIEW, src=(UPat.var("src"),), name="x"), - lambda x,src: isinstance(x.arg, ShapeTracker) and src.op is not Ops.STORE and x.dtype.base == src.dtype.base), - - (UPat(Ops.VALID, dtypes.bool, (UPat(Ops.VIEW),)), lambda: True), (UPat(Ops.CONST, src=(), name="x"), lambda x: type(x.arg) is type(dtypes.as_const(x.arg, x.dtype))), - # early LOAD has a - (UPat(Ops.LOAD, src=(UPat(Ops.VIEW, src=(UPat(GroupOp.Defines),)),)), lambda: True), - (UPat(Ops.LOAD, src=(UPat(Ops.VIEW, src=(UPat(GroupOp.Defines),)), UPat(Ops.STORE))), lambda: True), - - # early STORE has a - (UPat(Ops.STORE, src=(UPat(Ops.VIEW, src=(UPat(GroupOp.Defines),)), UPat())), lambda: True), - # **** new style load/store **** # make sure all index dtypes have been lowered @@ -243,20 +221,12 @@ spec = PatternMatcher([ # *** this is the UOp AST spec *** ast_spec = PatternMatcher([ - # VIEW can only exist in the edges - (UPat(Ops.VIEW, src=(UPat((Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL),))), lambda: True), - (UPat(Ops.VIEW, name="view"), lambda view: len(view.src) == 0), # all parent UOps must have the same shape (UPat(GroupOp.All-{Ops.SINK}, name="root"), lambda root: all_same([x.shape for x in root.src if x.st is not None])), ]) # *** this spec should match all UOps ever created *** -full_non_rangeify_spec = PatternMatcher([]) if RANGEIFY else PatternMatcher([ - # in non rangeify const can still have a View, and sometimes a FUSE while propagating - (UPat((Ops.VIEW, Ops.FUSE)).f(Ops.CONST), lambda: True), -]) - full_spec = PatternMatcher([ # SENTINEL should never be in the graph (UPat(Ops.SENTINEL), lambda: False), @@ -317,7 +287,7 @@ full_spec = PatternMatcher([ (UPat(Ops.DEFINE_VAR), lambda: True), # reshape on STORE (UPat(Ops.RESHAPE, src=(UPat(Ops.STORE),)), lambda: True), -])+full_non_rangeify_spec+tensor_uop_spec+spec +])+tensor_uop_spec+spec # ***** uop helpers ***** diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index 95e3f17a0d..9d2600c72e 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -16,7 +16,7 @@ from tinygrad.codegen.opt import axis_colors uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.VCONST: "#e0e0e0", Ops.REDUCE: "#FF5B5B", Ops.DEFINE_GLOBAL: "#ffe0b0", Ops.DEFINE_LOCAL: "#ffe0d0", Ops.DEFINE_REG: "#f0ffe0", Ops.REDUCE_AXIS: "#FF6B6B", Ops.RANGE: "#c8a0e0", Ops.ASSIGN: "#909090", Ops.BARRIER: "#ff8080", Ops.IF: "#c8b0c0", Ops.SPECIAL: "#c0c0ff", - Ops.INDEX: "#e8ffa0", Ops.WMMA: "#efefc0", Ops.VIEW: "#C8F9D4", Ops.MULTI: "#f6ccff", Ops.KERNEL: "#3e7f55", + Ops.INDEX: "#e8ffa0", Ops.WMMA: "#efefc0", Ops.MULTI: "#f6ccff", Ops.KERNEL: "#3e7f55", **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80", Ops.BUFFER_VIEW: "#E5EAFF", Ops.BLOCK: "#C4A484", Ops.BLOCKEND: "#C4A4A4", Ops.BUFFER: "#B0BDFF", Ops.COPY: "#a040a0", Ops.FUSE: "#FFa500", Ops.ALLREDUCE: "#ff40a0", Ops.MSELECT: "#d040a0", Ops.MSTACK: "#d040a0", Ops.CONTIGUOUS: "#FFC14D", Ops.REALIZE: "#C1C14D", @@ -63,14 +63,9 @@ def uop_to_json(x:UOp) -> dict[int, dict]: for u in (toposort:=x.toposort()): # always exclude DEVICE/CONST/UNIQUE if u.op in {Ops.DEVICE, Ops.CONST, Ops.UNIQUE} and u is not x: excluded.add(u) - # only exclude CONST VIEW source if it has no other children in the graph - if u.op is Ops.CONST and u.st is not None: excluded.update(u.src) for u in toposort: if u in excluded: continue argst = codecs.decode(str(u.arg), "unicode_escape") - if u.op is Ops.VIEW: - argst = ("\n".join([f"{shape_to_str(v.shape)} / {shape_to_str(v.strides)}"+("" if v.offset == 0 else f" / {srender(v.offset)}")+ - (f"\nMASK {mask_to_str(v.mask)}" if v.mask is not None else "") for v in unwrap(u.st).views])) if u.op in GroupOp.Movement: argst = (mask_to_str if u.op in {Ops.SHRINK, Ops.PAD} else shape_to_str)(u.arg) label = f"{str(u.op).split('.')[1]}{(chr(10)+word_wrap(argst.replace(':', ''))) if u.arg is not None else ''}" if u.dtype != dtypes.void: label += f"\n{u.dtype}" @@ -81,7 +76,7 @@ def uop_to_json(x:UOp) -> dict[int, dict]: try: if len(rngs:=u.ranges): label += f"\n({','.join([colored(range_str(x), axis_colors[x.arg[-1]]) for x in sorted(rngs, key=lambda x: x.arg[0:-1])])})" - if u.op not in {Ops.VIEW, Ops.BUFFER, Ops.KERNEL, Ops.ASSIGN, Ops.COPY, Ops.SINK, *GroupOp.Buffer} and u.st is not None: + if u.op not in {Ops.BUFFER, Ops.KERNEL, Ops.ASSIGN, Ops.COPY, Ops.SINK, *GroupOp.Buffer} and u.st is not None: label += f"\n{shape_to_str(u.shape)}" if u.op in {Ops.INDEX, Ops.BUFFERIZE}: label += f"\n{u.render()}"