From 127a7c8aee2d3554090c453fe5176c8c8f58347f Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Fri, 23 May 2025 15:27:09 +0300
Subject: [PATCH] assert AST views only exist in the edges (#10484)

* assert AST views only exist in the edges

* valid without device
---
 test/test_schedule.py        | 3 ++-
 test/unit/test_verify_ast.py | 5 ++---
 tinygrad/uop/spec.py         | 2 ++
 3 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/test/test_schedule.py b/test/test_schedule.py
index cc0508a0fa..336188de34 100644
--- a/test/test_schedule.py
+++ b/test/test_schedule.py
@@ -662,7 +662,8 @@ class TestSchedule(unittest.TestCase):
     a = Tensor.zeros(1, dtype=dtypes.int).contiguous().realize().lazydata
     b = Tensor.zeros(1, dtype=dtypes.int).contiguous().realize().lazydata
     c = Tensor.arange(4).realize().lazydata
-    kernel = UOp(Ops.KERNEL, src=(a, b, c), arg=Kernel(UOp.sink(c.r(Ops.ADD, (0,))+1, c.r(Ops.ADD, (0,))*2)))
+    kernel = UOp(Ops.KERNEL, src=(a, b, c.base), arg=Kernel(UOp.sink(c.r(Ops.ADD, (0,))+1, c.r(Ops.ADD, (0,))*2)))
+    assert all(s.op is Ops.BUFFER for s in kernel.src), f"views are not allowed here {kernel}"
     kernel = graph_rewrite(kernel, create_ast)
     run_schedule(check_schedule(UOp.sink(a.assign(kernel), b.assign(kernel)), 1))
     self.assertEqual(a.buffer.numpy(), [7])
diff --git a/test/unit/test_verify_ast.py b/test/unit/test_verify_ast.py
index eaa23ebcb9..74c1652f43 100644
--- a/test/unit/test_verify_ast.py
+++ b/test/unit/test_verify_ast.py
@@ -81,17 +81,16 @@ class TestVerifyAST(unittest.TestCase):
     const_st = [u.st for u in ast.toposort() if u.op is Ops.CONST][0]
     self.assertEqual(const_st, ShapeTracker.from_shape((1, 1)).expand((4, 4)))
 
-  @unittest.skip("questionable if we want this")
   def test_assert_swizzle(self):
     buf = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
     a = UOp(Ops.LOAD, dtypes.float, (buf, ShapeTracker.from_shape((32, 1)).to_uop()))
     r = UOp(Ops.REDUCE_AXIS, dtypes.float, (a,), (Ops.ADD, (0,)))
     st = UOp.store(buf, ShapeTracker.from_shape((32, 1)).to_uop(), r.view(r.st.expand((32, 1)))+a)
-    with self.assertRaisesRegex(InvalidASTException, "swizzle"): helper_test_verify_ast(st)
+    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, (UOp(Ops.DEVICE, arg="CPU"),), ShapeTracker.from_shape(())),))
+    a = UOp.const(dtypes.int, 0).replace(src=(UOp(Ops.VIEW, dtypes.void, (), ShapeTracker.from_shape(())),))
     st = UOp.store(buf, ShapeTracker.from_shape(()).to_uop(), a.cast(dtypes.float))
     helper_test_verify_ast(st)
 
diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py
index c4e2fd3386..68e42c102e 100644
--- a/tinygrad/uop/spec.py
+++ b/tinygrad/uop/spec.py
@@ -203,6 +203,8 @@ def verify_sink_dims(sink:UOp):
 shape_spec = PatternMatcher([
   # shapes must have either 1 or n in each dimension
   (UPat(Ops.SINK, src=UPat(Ops.STORE), name="sink"), verify_sink_dims),
+  # VIEW can only exist in the edges
+  (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])),
 ])