Revert "remove the on_stack stuff and any retries"

This reverts commit 49a2b328b9.

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)

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=(

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]]:

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),)))

test/test_rangeify.py

@@ -300,11 +300,12 @@ class TestOuterworld(unittest.TestCase):
     o.contiguous(i).realize()
     self.assertTrue((t==o).all().item())
 
-from tinygrad.schedule.rangeify import pm_rangeify, RangeifyContext
+@unittest.skip("pm_rangeify no longer exists. test this in a different way")
 class TestRangeifyPM(unittest.TestCase):
   def setUp(self): self.base = Tensor.empty(10*10).reshape(10, 10).contiguous()
   def assert_same(self, a, b):
     def run_pm_rangeify(t:Tensor):
+      from tinygrad.schedule.rangeify import pm_rangeify, RangeifyContext
       sink = t.uop.sink()
       pm_realize = PatternMatcher([(UPat(Ops.CONTIGUOUS, name="x"), lambda x: x.replace(op=Ops.REALIZE))])
       sink = graph_rewrite(sink, pm_realize)

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)

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)

test/unit/test_rewrite_not_ready.py

@@ -1,110 +0,0 @@
-import unittest
-from dataclasses import dataclass, field
-from tinygrad.uop.ops import PatternMatcher, UOp, graph_rewrite, Ops, UPat, GroupOp, RewriteNotReady
-
-# we could insert CHILDREN node
-
-@dataclass
-class ChildrenContext:
-  children: dict[UOp, list[UOp]]|None = None
-
-# this is a generic child labeller
-def extract_children(ctx:ChildrenContext, x:UOp):
-  if ctx.children is not None: return
-  ctx.children = {k:list(v.keys()) for k,v in x.get_consumer_map().items() if len(v) > 1}
-
-def mark_children(ctx:ChildrenContext, x:UOp):
-  new_srcs = [(UOp(Ops.CHILD, s.dtype, src=(s,), arg=(ctx.children[s].index(x), len(ctx.children[s]))) if s in ctx.children else s) for s in x.src]
-  return x.replace(src=tuple(new_srcs))
-
-pm_children = PatternMatcher([
-  (UPat(Ops.SINK, name="x"), extract_children),
-  (UPat(GroupOp.All-{Ops.CHILD}, name="x"), mark_children),
-])
-
-@dataclass
-class TestContext:
-  seen_children: dict[UOp, set[int]] = field(default_factory=dict)
-  ready_children: dict[UOp, set[int]] = field(default_factory=dict)
-  seen_consts:int = 0
-  saved_seen_consts:int = 0
-  exp2_visit_count:int = 0
-
-# this is a generic pattern
-def visit_child(ctx:ChildrenContext, x:UOp):
-  if x.src[0] not in ctx.seen_children:
-    ctx.seen_children[x.src[0]] = set()
-    ctx.ready_children[x.src[0]] = set()
-  ctx.seen_children[x.src[0]].add(x.arg[0])
-  if len(ctx.seen_children[x.src[0]]) != x.arg[1]:
-    print(f"visit CHILD {x.arg} bottom up -- not ready {ctx.seen_children[x.src[0]]}")
-    raise RewriteNotReady
-  print(f"visit CHILD {x.arg} bottom up -- READY {ctx.seen_children[x.src[0]]}")
-  ctx.ready_children[x.src[0]].add(x.arg[0])
-
-pm_child_visitor = PatternMatcher([
-  (UPat(Ops.CHILD, name="x"), visit_child),
-])
-
-# this is for the test
-def see_const(ctx:ChildrenContext, c:UOp): ctx.seen_consts += c.arg
-def see_exp2(ctx:ChildrenContext): ctx.exp2_visit_count += 1
-def save_seen_consts(ctx:ChildrenContext, x:UOp): ctx.saved_seen_consts = ctx.seen_consts
-pm_consts = PatternMatcher([
-  (UPat(Ops.DEFINE_VAR, name="x"), save_seen_consts),
-  (UPat()+UPat.cvar("c"), see_const),
-  (UPat(Ops.EXP2), see_exp2),
-])
-
-class TestChildrenRewrite(unittest.TestCase):
-  def test_not_ready_double_simple(self):
-    global_a = UOp.variable("a", 0, 10).exp2()
-    inter = (global_a+global_a).exp2()
-    global_sink = (inter+inter).sink()
-
-    sink = graph_rewrite(global_sink, pm_children, ctx=ChildrenContext(), bottom_up=True)
-    ctx = TestContext()
-    graph_rewrite(sink, pm_consts, ctx=ctx, bottom_up=True)
-    self.assertEqual(ctx.exp2_visit_count, 2)
-
-  def test_not_ready_double(self):
-    global_a = UOp.variable("a", 0, 10).exp2()
-    inter = ((global_a+1000)+(global_a+100)).exp2()
-    global_sink = ((inter+10)+(inter+1)).sink()
-
-    sink = graph_rewrite(global_sink, pm_children, ctx=ChildrenContext(), bottom_up=True)
-    print("test_not_ready_double")
-    ctx = TestContext()
-    graph_rewrite(sink, pm_child_visitor+pm_consts, ctx=ctx, bottom_up=True)
-    self.assertEqual(ctx.exp2_visit_count, 2)
-    self.assertEqual(ctx.seen_consts, ctx.saved_seen_consts)
-    self.assertEqual(ctx.seen_consts, 1111)
-
-  def test_in_srcs_twice(self):
-    global_a = UOp.variable("a", 0, 10).exp2()
-    global_sink = (global_a+global_a).sink()
-
-    ctx = TestContext()
-    graph_rewrite(global_sink, pm_consts, ctx=ctx, bottom_up=True)
-    self.assertEqual(ctx.exp2_visit_count, 1)
-
-  def test_not_ready(self):
-    global_a = UOp.variable("a", 0, 10).exp2()
-    global_sink = ((global_a+2)+(global_a+3)).sink()
-
-    # without children and not ready, we don't see both adds before the DEFINE_VAR
-    ctx = TestContext()
-    graph_rewrite(global_sink, pm_consts, ctx=ctx, bottom_up=True)
-    self.assertNotEqual(ctx.seen_consts, ctx.saved_seen_consts)
-    self.assertEqual(ctx.exp2_visit_count, 1)
-
-    # with children and not ready we do
-    sink = graph_rewrite(global_sink, pm_children, ctx=ChildrenContext(), bottom_up=True)
-    ctx = TestContext()
-    graph_rewrite(sink, pm_child_visitor+pm_consts, ctx=ctx, bottom_up=True)
-    self.assertEqual(ctx.seen_consts, ctx.saved_seen_consts)
-    self.assertEqual(ctx.exp2_visit_count, 1)
-    self.assertSetEqual(list(ctx.ready_children.values())[0], {0,1})
-
-if __name__ == '__main__':
-  unittest.main()

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)
 

test/unit/test_uop_spec.py

@@ -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()

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

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,))),
-])
+])

tinygrad/schedule/indexing.py

@@ -7,6 +7,32 @@ from tinygrad.uop.symbolic import sym
 from tinygrad.helpers import argsort, all_same, Context
 from tinygrad.uop.ops import graph_rewrite, sint, AxisType
 
+ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.ASSIGN, Ops.COPY, Ops.BUFFER, Ops.BUFFER_VIEW,
+                     Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.DEFINE_GLOBAL,
+                     Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.LOAD, Ops.KERNEL}
+
+def realize(ctx:dict[UOp, None], tr:UOp) -> None: ctx[tr] = None
+
+def realize_srcs(ctx:dict[UOp, None], rb:UOp) -> None:
+  for s in rb.src:
+    if s.base.op not in ALWAYS_CONTIGUOUS: ctx[s] = None
+
+def realize_assign(ctx:dict[UOp, None], a:UOp) -> None:
+  if a.src[1].op not in ALWAYS_CONTIGUOUS: ctx[a.src[1]] = None
+  # if it's a kernel, we don't realize it
+  if a.src[1].op is not Ops.KERNEL: ctx[a] = None
+
+pm_generate_realize_map = PatternMatcher([
+  # always realize SINK src
+  (UPat(Ops.SINK, name="s"), lambda ctx,s: ctx.update((x.base, None) for x in s.src if x.base.op not in ALWAYS_CONTIGUOUS)),
+  # always realize ASSIGN/COPY/BUFFER_VIEW/CONTIGUOUS
+  (UPat({Ops.COPY, Ops.BUFFER_VIEW, Ops.CONTIGUOUS}, name="tr"), realize),
+  # realize srcs of COPY, MSELECT, MSTACK
+  (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK), name="rb"), realize_srcs),
+  # realize input to assign (might be optimized out)
+  (UPat(Ops.ASSIGN, name="a"), realize_assign),
+])
+
 @dataclass(frozen=True)
 class BufferizeOpts:
   # on AddrSpace.LOCAL, device is the id
@@ -77,11 +103,15 @@ pm_apply_rangeify = PatternMatcher([
   (UPat((Ops.CONST, Ops.DEFINE_VAR), name="c"), lambda ctx,c: c.replace(src=()) if c in ctx.range_map else None),
 ])
 
-def run_rangeify(tsink:UOp, realize_map:dict[UOp, None], debug) -> tuple[UOp, IndexingContext]:
+def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
   tsink_base = UOp.sink(*[x.base for x in tsink.src])
 
-  # explicit rangeify
   rctx = IndexingContext()
+
+  # get ops to realize
+  graph_rewrite(tsink, pm_generate_realize_map, ctx=rctx.realize_map, name="Input Graph")
+
+  # explicit rangeify
   ending_ranges: dict[UOp, bool] = {}
   for x in tsink_base.reverse_toposort(consumer_map:=tsink_base.get_consumer_map()):
     if x.op in {Ops.DEVICE, Ops.UNIQUE}: continue
@@ -89,9 +119,9 @@ def run_rangeify(tsink:UOp, realize_map:dict[UOp, None], debug) -> tuple[UOp, In
 
     # if this element has weight and it's ending a range, we (force) realize it
     if ending_ranges[x] and x.op in GroupOp.Elementwise.union({Ops.REDUCE_AXIS}):
-      if x.op_in_backward_slice_with_self(Ops.BUFFER, Ops.REALIZE, Ops.BUFFERIZE, Ops.CONTIGUOUS):
+      if x.op_in_backward_slice_with_self(Ops.BUFFER, Ops.BUFFERIZE, Ops.CONTIGUOUS):
         if x.op_in_backward_slice_with_self(Ops.REDUCE_AXIS):
-          realize_map[x] = None
+          rctx.realize_map[x] = None
 
     # *** the ranges on the output are
     #  1. new if this op is realized
@@ -99,7 +129,7 @@ def run_rangeify(tsink:UOp, realize_map:dict[UOp, None], debug) -> tuple[UOp, In
     #  3. potentially new if this op has 2+ consumers
 
     consumer_rngs = [rctx.range_map[c][0] for c in consumer_map[x] if c in rctx.range_map]
-    if x in realize_map:
+    if x in rctx.realize_map:
       # if this is in the realize_map, we create new ranges (at the output)
       out_rngs = [rctx.new_range(s) for s in x.shape]
       # all ranges are ended now
@@ -136,7 +166,7 @@ def run_rangeify(tsink:UOp, realize_map:dict[UOp, None], debug) -> tuple[UOp, In
           out_rngs.append(rctx.new_range(x.shape[i]))
 
       # we have to realize here if there's new ranges
-      if not all_all_same: realize_map[x] = None
+      if not all_all_same: rctx.realize_map[x] = None
 
     # TODO: some ops don't have shape, enable this after the `.st` property is removed
     #assert len(out_rngs) == len(x.shape), \
@@ -190,12 +220,11 @@ def run_rangeify(tsink:UOp, realize_map:dict[UOp, None], debug) -> tuple[UOp, In
         if i in x.arg[1]: rngs[i] = rctx.new_range(s, axistype=AxisType.REDUCE)
 
     if debug:
-      print("***" if x in realize_map else "   ", len(consumer_map[x]), f"{str(x.op):20s}",
+      print("***" if x in rctx.realize_map else "   ", len(consumer_map[x]), f"{str(x.op):20s}",
             UOp.sink().index(*rngs).render(), " -> ", UOp.sink().index(*out_rngs).render())
 
     # assign to the range map. rngs are the input ranges, out_rngs are the output ranges, from the x op.
     rctx.range_map[x] = (rngs, out_rngs)
 
-  rctx.realize_map = realize_map
   tsink = graph_rewrite(tsink, pm_apply_rangeify, ctx=rctx, bottom_up=True, name="apply rangeify")
   return tsink, rctx

tinygrad/schedule/rangeify.py

@@ -1,16 +1,14 @@
-from typing import Any, cast, Iterator
-
-import functools, operator, itertools
+from typing import cast
 from dataclasses import dataclass, field
 from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, RewriteNotReady, _substitute, ssimplify, KernelInfo
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, _substitute, ssimplify, KernelInfo
 from tinygrad.uop.symbolic import sym, symbolic_simple
 from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, RANGEIFY, Context, flatten, dedup, unwrap, all_int, DEBUG, SPLIT_REDUCEOP
 from tinygrad.helpers import Metadata
 from tinygrad.uop.ops import track_rewrites, graph_rewrite, identity_element, sint, AxisType
 from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_unparented
 from tinygrad.codegen.opt import Opt
-from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, IndexingContext
+from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, ALWAYS_CONTIGUOUS, IndexingContext
 
 # creation can recurse a lot
 import sys
@@ -19,10 +17,6 @@ sys.setrecursionlimit(10000)
 # *****************
 # 0. do some cleanup rewrites, mostly copied from the old stuff
 
-ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.ASSIGN, Ops.COPY, Ops.BUFFER, Ops.BUFFER_VIEW,
-                     Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.DEFINE_GLOBAL,
-                     Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.LOAD, Ops.KERNEL}
-
 def find_permutes(a:UOp, b:UOp, assign:UOp):
   if not (permutes:=[s for s in b.toposort(gate=lambda s:s.op not in ALWAYS_CONTIGUOUS)
                      if s.op in GroupOp.Movement and s.op not in {Ops.RESHAPE, Ops.EXPAND, Ops.PAD, Ops.SHRINK}]): return
@@ -104,79 +98,9 @@ earliest_rewrites = PatternMatcher([
   (UPat(Ops.CONTIGUOUS, name="root", src=(UPat(Ops.BUFFER),)), lambda root: root.src[0].forced_reshape(root.shape).rtag(root.tag)),
 ])
 
-# *****************
-# 1. add realize where we have to
-
-def realize(ctx:dict[UOp, None], tr:UOp) -> None: ctx[tr] = None
-
-def realize_parents(ctx:dict[UOp, None], rb:UOp) -> None:
-  for s in rb.src:
-    if s.base.op not in ALWAYS_CONTIGUOUS: ctx[s] = None
-
-def realize_assign(ctx:dict[UOp, None], a:UOp) -> None:
-  if a.src[1].op not in ALWAYS_CONTIGUOUS: ctx[a.src[1]] = None
-  # if it's a kernel, we don't realize it
-  if a.src[1].op is not Ops.KERNEL: ctx[a] = None
-
-do_realize = PatternMatcher([
-  # always realize SINK parents
-  (UPat(Ops.SINK, name="s"), lambda ctx,s: ctx.update((x.base, None) for x in s.src if x.base.op not in ALWAYS_CONTIGUOUS)),
-  # always realize ASSIGN/COPY/BUFFER_VIEW/CONTIGUOUS
-  (UPat({Ops.COPY, Ops.BUFFER_VIEW, Ops.CONTIGUOUS}, name="tr"), realize),
-  # realize parents of COPY, MSELECT, MSTACK
-  (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK), name="rb"), realize_parents),
-  # realize input to assign (might be optimized out)
-  (UPat(Ops.ASSIGN, name="a"), realize_assign),
-])
-
-class WrappedContig:
-  def __init__(self, x): self.x = x
-  def __repr__(self): return f"C({self.x})"
-add_contiguous = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda ctx,x: x.replace(tag=WrappedContig(x.tag)).realize() if x in ctx else None),])
-remove_contig_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=x.tag.x) if isinstance(x.tag, WrappedContig) else None)])
-
-# *****************
-# 2. mark all children
-
-@dataclass
-class ChildrenContext: children: dict[UOp, list[UOp]]|None = None
-def extract_children(ctx:ChildrenContext, x:UOp):
-  if ctx.children is not None: return
-  children_map = x.get_consumer_map()
-  ctx.children = {}
-  for k,v in children_map.items():
-    # NOTE: we treat mstack children like sink here
-    non_sink_children = [u for u in v if u.op not in {Ops.SINK, Ops.MSTACK}]
-    if len(non_sink_children) <= 1: continue
-    # NOTE: this gate shouldn't be here
-    if k.op_in_backward_slice_with_self(Ops.REDUCE_AXIS) and k.op_in_backward_slice_with_self(Ops.BUFFER, Ops.CONTIGUOUS):
-      ctx.children[k] = non_sink_children
-
-def mark_children(ctx:ChildrenContext, x:UOp):
-  assert ctx.children is not None
-  new_srcs = [(UOp(Ops.CHILD, s.dtype, src=(UOp(Ops.CHILDREN, s.dtype, (s,), arg=len(ctx.children[s])),),
-                   arg=(ctx.children[s].index(x), len(ctx.children[s]))) if s in ctx.children else s) for s in x.src]
-  return x.replace(src=tuple(new_srcs))
-
-pm_children = PatternMatcher([
-  (UPat(Ops.SINK, name="x"), extract_children),
-  (UPat(GroupOp.All-{Ops.CHILD, Ops.CHILDREN, Ops.SINK}, name="x"), mark_children),
-])
-
 # *****************
 # 3a. rangeify (movement)
-
-@dataclass
-class RangeifyContext:
-  # block on parent until all children have been seen
-  seen_children: dict[UOp, dict[int, UOp]] = field(default_factory=dict)
-  seen_child: dict[UOp, Any] = field(default_factory=dict)
-  progress: int = 0
-
-  # create ranges
-  range_idx: Iterator[int] = field(default_factory=itertools.count)
-  def new_range(self, s:sint, axistype:AxisType=AxisType.LOOP):
-    return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(dtypes.index, 0)
+# NOTE: this can be deleted after the cleanup is refactored
 
 def map_reshape(idx:UOp, r:UOp):
   acc = 1
@@ -241,150 +165,6 @@ pm_mops = PatternMatcher([
   (UPat(Ops.PAD, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"), map_pad),
 ])
 
-# *****************
-# 3b. rangeify (ops)
-
-# bufferization can happen in three ways
-#   1. there's an explicit REALIZE in the graph
-#   2. the ranges from the children don't match and we have to create a buffer (only on children)
-#   3. might_end_axis triggers because we should be closing a loop to save compute
-
-def map_partial_realize(ctx:RangeifyContext, x:UOp, idx:UOp):
-  if x.arg is None: return None  # map_contiguous can handle this
-  # NOTE: all partial contiguous can safely be replaced by full contiguous. we should be able to match old functionality like this
-  if not (RANGEIFY > 1): return idx.replace(src=(x.replace(arg=None),)+idx.src[1:])
-  ranges = []
-  new_ranges = []
-  passthrough_idx = []
-  for i,s in enumerate(x.shape):
-    if i not in x.arg:
-      ranges.append(idx.src[1+i])
-      continue
-    passthrough_idx.append(idx.src[1+i])
-    ranges.append(ctx.new_range(s))
-    new_ranges.append(ranges[-1])
-  # TODO: this should be able to be global or local
-  ret = x.src[0].index(*ranges).bufferize(*[x for x in new_ranges if x.op is not Ops.CONST],
-                                          arg=BufferizeOpts(device=None, addrspace=AddrSpace.LOCAL))
-  return ret.index(*passthrough_idx)
-
-def map_realize(ctx:RangeifyContext, x:UOp):
-  if x.arg is not None: return None
-  ranges = [ctx.new_range(s) for s in x.shape]
-  return x.src[0].index(*ranges).bufferize(*x.src[1:], *ranges, arg=BufferizeOpts(device=x.device), tag=x.src[0].tag)
-
-def map_reduce(ctx:RangeifyContext, idx:UOp, red:UOp):
-  rngs = list(idx.src[1:])
-  new_ranges = []
-  for i,s in enumerate(red.src[0].shape):
-    if i in red.arg[1]:
-      rngs[i] = ctx.new_range(s, axistype=AxisType.REDUCE)
-      new_ranges.append(rngs[i])
-  return UOp(Ops.REDUCE, red.dtype, src=(red.src[0].index(*rngs),)+tuple(new_ranges), arg=red.arg[0], tag=red.tag)
-
-def index_child(ctx:RangeifyContext, c:UOp, x:UOp, idx:UOp):
-  if c not in ctx.seen_children: ctx.seen_children[c] = {}
-  ctx.seen_children[c][x.arg[0]] = idx
-  # wait here until we have seen all the children
-  if len(ctx.seen_children[c]) != x.arg[1]:
-    ctx.progress += 1
-    if ctx.progress > 10000: raise RuntimeError("children not making progress")
-    raise RewriteNotReady
-  ctx.progress = 0
-
-  if c not in ctx.seen_child:
-    all_rngs = list(zip(*[ch.src[1:] for ch in ctx.seen_children[c].values()]))
-    out_rngs = []
-    end_ranges = []
-    idx_ranges = []
-    # NOTE: locals aren't working, so we only fully bufferize here (unless RANGEIFY > 1)
-    rngs_valids = []
-    for valid_rngs in all_rngs:
-      rngs, valids = zip(*[(r.get_idx(), r.get_valid()) for r in valid_rngs])
-      # if a range has a 1 src, it's the same as UOp.const(dtypes.index, 0)
-      same_rngs = [x if x.op is not Ops.RANGE or resolve(x.src[0] != 1) else UOp.const(dtypes.index, 0) for x in rngs]
-      rngs_valids.append((rngs, valids, all_same(same_rngs)))
-    all_all_same = all(same_rngs for _,_,same_rngs in rngs_valids)
-    for i,(rngs,valids,same_rngs) in enumerate(rngs_valids):
-      # we compare the ranges without their valids
-      if same_rngs and (all_all_same or RANGEIFY > 1):
-        # the new valid is the OR of all the children valids
-        minimum_valid = functools.reduce(operator.or_, valids, UOp.const(dtypes.bool, False))
-        out_rngs.append(minimum_valid.where(rngs[0], UOp.invalid()).simplify())
-      else:
-        out_rngs.append(ctx.new_range(c.shape[i]))
-        end_ranges.append(out_rngs[-1])
-        idx_ranges.append(i)
-    ctx.seen_child[c] = (out_rngs, idx_ranges, end_ranges)
-  else:
-    out_rngs, idx_ranges, end_ranges = ctx.seen_child[c]
-    for i,nr in zip(idx_ranges, end_ranges): out_rngs[i] = nr
-  # index based on the shared ranges
-  ret = c.index(*out_rngs)
-  # if all ranges aren't the same between children, we have to bufferize
-  if len(idx_ranges) > 0:
-    if len(idx_ranges) == len(out_rngs):
-      # this is a global bufferize
-      ret = ret.bufferize(*end_ranges, arg=BufferizeOpts(device=x.device))
-    else:
-      assert RANGEIFY > 1, "this isn't supported with RANGEIFY=1"
-      ret = ret.bufferize(*end_ranges, arg=BufferizeOpts(device=None, addrspace=AddrSpace.LOCAL))
-    ret = ret.index(*[idx.src[1+i] for i in idx_ranges])
-  return ret
-
-def children_gate(ctx:RangeifyContext, idx:UOp, c:UOp):
-  if len(ctx.seen_children[c]) != c.arg: raise RuntimeError("all children should have been seen by now")
-  return idx.replace(src=(idx.src[0].src[0],)+idx.src[1:])
-
-def might_end_axis(idx:UOp):
-  if idx.arg is None: return None
-  # TODO: write a proper cost function here
-  if not idx.op_in_backward_slice_with_self(Ops.BUFFER, Ops.REALIZE, Ops.BUFFERIZE): return None
-  if not idx.op_in_backward_slice_with_self(Ops.REDUCE_AXIS): return None
-  to_end_axis = []
-  for i,a in enumerate(idx.src[1:]):
-    # in RANGEIFY=1, always realize
-    if not (RANGEIFY > 1) or any(x.arg > idx.arg for x in a.toposort() if x.op is Ops.RANGE):
-      to_end_axis.append(i)
-  if to_end_axis: return idx.replace(src=(idx.src[0].realize(arg=tuple(to_end_axis)),)+idx.src[1:], arg=None)
-  return idx.replace(arg=None)
-
-def unprocessed_index(x:UOp): raise RuntimeError(f"unprocessed index on {x.src[0].op}")
-
-pm_rangeify = pm_mops+PatternMatcher([
-  # sink contigs to kick it off
-  (UPat(Ops.REALIZE, src=(UPat(),), name="x", allow_any_len=True), map_realize),
-  # if there's an INDEX it can support partial contig
-  (UPat(Ops.INDEX, src=(UPat(Ops.REALIZE, src=(UPat(),), name="x"),), allow_any_len=True, name="idx"), map_partial_realize),
-
-  # if there are new ended children, tag the SINK
-  (UPat(Ops.INDEX, src=(UPat(Ops.CHILD, src=(UPat(name="c"), ), name="x"),), allow_any_len=True, name="idx"), index_child),
-  (UPat(Ops.INDEX, src=(UPat(Ops.CHILDREN, name="c"),), allow_any_len=True, name="idx"), children_gate),
-
-  # if we come across this, remove it. it was a CHILD unused in an INDEX
-  (UPat(Ops.CHILD, src=(UPat(Ops.CHILDREN, src=(UPat.var("x"),)),)), lambda x: x),
-
-  # CONST (or DEFINE_VAR) can't have axes. remove INDEX when we get here
-  (UPat(Ops.INDEX, src=(UPat((Ops.CONST, Ops.DEFINE_VAR), name="c"),)), lambda c: c.replace(src=())),
-
-  # handle arg on any op with weight. old endrange stuff
-  (UPat(Ops.INDEX, src=(UPat(GroupOp.Elementwise.union({Ops.REDUCE_AXIS})),), allow_any_len=True, name="idx"), might_end_axis),
-
-  # handle assign
-  (UPat(Ops.INDEX, src=(UPat(Ops.ASSIGN, name="assign"),), allow_any_len=True, name="x"),
-    lambda x,assign: assign.replace(src=tuple([s.index(*x.src[1:]) for s in assign.src])+(assign.src[0],)) \
-        if assign.src[1].op is not Ops.KERNEL else None),
-
-  # move MAP through elementwise ALU / reduce. these are the items with cost
-  (UPat(Ops.INDEX, src=(UPat(GroupOp.Elementwise.union(
-    {Ops.STORE, Ops.COPY, Ops.BUFFER_VIEW, Ops.DEVICE, Ops.BIND, Ops.CONTIGUOUS, Ops.NOOP})),), allow_any_len=True, name="x"),
-   lambda x: x.src[0].replace(src=tuple([s.index(*x.src[1:]) for s in x.src[0].src]))),
-  (UPat(Ops.INDEX, src=(UPat(Ops.REDUCE_AXIS, name="red"),), allow_any_len=True, name="idx"), map_reduce),
-
-  # assert if there's any index we didn't process
-  (UPat(GroupOp.All-{Ops.REALIZE, Ops.BUFFERIZE, Ops.MSELECT, Ops.MSTACK}).f(Ops.INDEX, name="x"), unprocessed_index),
-])
-
 # *****************
 # 3.5 cleanups
 
@@ -500,7 +280,7 @@ to_bufferview = PatternMatcher([
 ])
 
 DEVICE_MAX_BUFS = {"METAL": 31, "WEBGPU": 8} # TODO: get from device?
-def limit_bufs(ctx:RangeifyContext, root:UOp):
+def limit_bufs(ctx:IndexingContext, root:UOp):
   if (device:=root._device) is None: return None # no device, index related calculations
   device = device if isinstance(device, str) else device[0].split(":")[0]
   if not (MAX_BUFS:=getenv("MAX_KERNEL_BUFFERS", DEVICE_MAX_BUFS.get(device, 0))): return None
@@ -760,19 +540,9 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
   tsink = graph_rewrite(sink, add_tags, ctx=uop_list, bottom_up=True, name="number the uops")
 
   tsink = graph_rewrite(tsink, earliest_rewrites+replace_contiguous, ctx={}, name="earliest rewrites")
-  realize_map: dict[UOp, None] = {}
-  graph_rewrite(tsink, do_realize, ctx=realize_map, name="Input Graph")
 
-  FAST = getenv("FAST", 1)
-  if FAST:
-    rctx: RangeifyContext|IndexingContext
-    tsink, rctx = run_rangeify(tsink, realize_map, FAST > 1)
-  else:
-    # NOTE: we don't use contiguous here, contiguous is a user op
-    tsink = graph_rewrite(tsink, add_contiguous, ctx=realize_map, bottom_up=True, name="add realize")
-    tsink = graph_rewrite(tsink, remove_contig_tags, name="remove contiguous tags")
-    tsink = graph_rewrite(tsink, pm_children, ctx=ChildrenContext(), bottom_up=True, name="get children")
-    tsink = graph_rewrite(tsink, pm_rangeify, ctx=(rctx:=RangeifyContext()), bottom_up=True, name="rangeify")
+  # convert movement ops to ranges
+  tsink, rctx = run_rangeify(tsink, getenv("DEBUG_RANGEIFY", 0))
 
   # NOTE: sym (vs symbolic_simple) breaks things here because ranges with len 1 aren't handled right
   tsink = graph_rewrite(tsink, symbolic_simple+pm_reduce_unparented, name="symbolic")  # this supports const folding

tinygrad/uop/__init__.py

@@ -12,9 +12,6 @@ class Ops(FastEnum):
   NOOP = auto(); SINK = auto(); UNIQUE = auto(); DEVICE = auto(); KERNEL = auto(); PRECAST = auto(); REWRITE_ERROR = auto()  # noqa: E702
   SENTINEL = auto()
 
-  # track children
-  CHILD = auto(); CHILDREN = auto() # noqa: E702
-
   # buffer ops
   COPY = auto(); BUFFER = auto(); BUFFER_VIEW = auto(); MSELECT = auto(); MSTACK = auto() # noqa: E702
 
@@ -24,7 +21,6 @@ class Ops(FastEnum):
 
   # ops that adjust the behavior of the scheduler
   CONTIGUOUS = auto(); CONTIGUOUS_BACKWARD = auto(); DETACH = auto(); FUSE = auto() # noqa: E702
-  REALIZE = auto()
 
   # blocks in linearizer (only used there)
   BLOCK = auto(); BLOCKSTART = auto(); BLOCKEND = auto(); BLOCKFINAL = auto() # noqa: E702
@@ -33,12 +29,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 +90,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

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):
@@ -384,7 +367,6 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     return self.src[0] if self.op is Ops.WHERE and self.src[2].arg is Invalid else UOp.const(dtypes.bool, self.arg is not Invalid)
   def reduce(self, *src:UOp, **kwargs): return UOp(Ops.REDUCE, kwargs.pop('dtype', self.dtype), src=(self,)+src, **kwargs)
   def contiguous(self, *args, **kwargs): return UOp(Ops.CONTIGUOUS, dtype=self.dtype, src=(self,)+args, **kwargs)
-  def realize(self, *args, **kwargs): return UOp(Ops.REALIZE, dtype=self.dtype, src=(self,)+args, **kwargs)
   def contiguous_backward(self): return self.alu(Ops.CONTIGUOUS_BACKWARD)
   def bufferize(self, *args, **kwargs): return UOp(Ops.BUFFERIZE, dtype=self.dtype, src=(self,)+args, **kwargs)
   def fuse(self): return self.alu(Ops.FUSE)
@@ -456,10 +438,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 +553,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 +770,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)
@@ -974,8 +953,8 @@ class TrackedPatternMatcher(PatternMatcher):
         continue
       match_stats[p][1] += 1
       try: ret = match(uop, ctx)
-      except Exception as e:
-        if TRACK_MATCH_STATS >= 2 and active_rewrites and not isinstance(e, RewriteNotReady):
+      except Exception:
+        if TRACK_MATCH_STATS >= 2 and active_rewrites:
           active_rewrites[-1].matches.append((track_uop(uop), track_uop(UOp(Ops.REWRITE_ERROR, src=uop.src, arg=str(sys.exc_info()[1]))), p.location))
         raise
       if ret is not None and ret is not uop:
@@ -1018,7 +997,6 @@ if TRACK_MATCH_STATS or PROFILE:
 # *** simple graph rewrite engine ***
 
 with Context(SPEC=0): SENTINEL = UOp(Ops.SENTINEL)
-class RewriteNotReady(Exception): pass
 class BottomUpGate(Exception): pass
 class RewriteContext:
   def __init__(self, pm, bpm, ctx=None):
@@ -1058,10 +1036,6 @@ class RewriteContext:
               if test_n in seen: raise RuntimeError("infinite loop in fixed_point_rewrite")
               seen.add(test_n)
               new_n, test_n = test_n, self.cached_bpm_rewrite(test_n)
-          except RewriteNotReady:
-            # try the full thing again later
-            stack.appendleft((n, 0, n))
-            continue
           except BottomUpGate:
             # if the bpm matching raised a gate, we are done with this node and dont continue down the srcs
             self.replace[n] = new_n
@@ -1075,7 +1049,7 @@ class RewriteContext:
         tmp = []
         for x in new_n.src:
           if (rx:=self.replace.get(x, SENTINEL)) is SENTINEL:
-            # if some new sources aren't ready, we try this again later
+            # if some new sources aren't ready, we try this again later. happens with on_stack, maybe should remove?
             stack.appendleft((n, 1, new_n))
             break
           tmp.append(rx)
@@ -1168,7 +1142,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 +1168,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 +1176,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] = []

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 <bufview, store?>
-  (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 <bufview, val>
-  (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),
@@ -269,11 +239,6 @@ full_spec = PatternMatcher([
   # where on index in rhs position is fine
   (UPat(Ops.WHERE, src=(UPat(dtype=dtypes.bool), UPat(), UPat(dtype=dtypes.index))), lambda: True),
 
-  # all children is fine
-  (UPat(Ops.CHILDREN), lambda: True),
-  # child must have CHILDREN parent
-  (UPat(Ops.CHILD, src=(UPat(Ops.CHILDREN),)), lambda: True),
-
   # all rewrite error are okay
   (UPat(Ops.REWRITE_ERROR), lambda: True),
 
@@ -281,8 +246,6 @@ full_spec = PatternMatcher([
   (UPat(Ops.BUFFER_VIEW, src=(UPat((Ops.INDEX, Ops.LOAD)),)), lambda: True),
   # bufferize (must be on ranges)
   (UPat(Ops.BUFFERIZE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.op in {Ops.RANGE, Ops.CONST} for y in x.src[1:])),
-  # realize with one src is fine
-  (UPat(Ops.REALIZE, src=(UPat(),)), lambda: True),
   # intermediate index
   (UPat(Ops.INDEX, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.dtype == dtypes.index for y in x.src[1:]) or None),
   (UPat(Ops.REDUCE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.dtype == dtypes.index for y in x.src[1:])),
@@ -317,7 +280,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 *****
 

tinygrad/viz/serve.py

@@ -16,12 +16,11 @@ 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",
-               Ops.CHILDREN: "#80ffc0", Ops.CHILD: "#80fff0", Ops.BUFFERIZE: "#FF991C", Ops.REWRITE_ERROR: "#ff2e2e",
-               Ops.SUBSTITUTE: "#ffff00"}
+               Ops.ALLREDUCE: "#ff40a0", Ops.MSELECT: "#d040a0", Ops.MSTACK: "#d040a0", Ops.CONTIGUOUS: "#FFC14D",
+               Ops.BUFFERIZE: "#FF991C", Ops.REWRITE_ERROR: "#ff2e2e", Ops.SUBSTITUTE: "#ffff00"}
 
 # VIZ API
 
@@ -63,14 +62,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 +75,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()}"