more test relaxations from prealloc_bufs (#14880)

test/backend/test_schedule.py

@@ -168,13 +168,13 @@ class TestSchedule(unittest.TestCase):
     a = Tensor.full((4,), 4.0).contiguous().realize()
     b = Tensor.full((4,), 2.0).contiguous().realize()
     expr = (a*b)/b
-    check_schedule(expr, 0)
+    run_schedule(check_schedule(expr, 0))
     np.testing.assert_allclose(expr.numpy(), np.full((4,), 4.0))
 
   def test_div_collapse_const(self):
     a = Tensor.full((4,), 4.0).contiguous().realize()
     expr = a/a
-    check_schedule(expr, 0)
+    run_schedule(check_schedule(expr, 0))
     np.testing.assert_allclose(expr.numpy(), np.full((4,), 1.0))
 
   def test_div_collapse(self):
@@ -1236,11 +1236,12 @@ class TestView(unittest.TestCase):
     bv = b.pad(((0, 2),))[-2:]
     # this becomes a late a*0
     late_mul = a*bv
-    check_schedule(late_mul, 0)
+    run_schedule(check_schedule(late_mul, 0))
+    # NOTE: no longer checked
     # the arange doesn't realize
-    self.assertIsNone(b.uop.base.realized)
+    #self.assertIsNone(b.uop.base.realized)
     # mul doesn't realize
-    self.assertIsNone(late_mul.uop.base.realized)
+    #self.assertIsNone(late_mul.uop.base.realized)
     self.assertEqual(late_mul.tolist(), [0, 0])
 
   # SINK has two branches:
@@ -1256,17 +1257,18 @@ class TestView(unittest.TestCase):
     s = check_schedule([late_mul, other_child], 2)
     # the arange becomes a BUFFER
     self.assertIs(b.uop.base.op, Ops.BUFFER)
+    # NOTE: no longer checked
     # mul still collapses
-    self.assertIs(late_mul.uop.base.op, Ops.CONST)
+    #self.assertIs(late_mul.uop.base.op, Ops.CONST)
     run_schedule(s)
     self.assertEqual(other_child.tolist(), [2, 3, 4])
 
 @unittest.skipIf(Device.DEFAULT == "CPU", "tests copy from another device to cpu")
 class TestCopyFolding(unittest.TestCase):
   def test_const_copy_is_free(self):
-    b = Tensor(1).to("CPU")
-    check_schedule(b, 0, filter_sink=False)
-    assert b.item() == 1
+    b = Tensor(1).to("CPU") * 4
+    run_schedule(check_schedule(b, 0, filter_sink=False))
+    assert b.item() == 4
 
   def test_one_hot_with_copy(self):
     y = Tensor([1, 2, 3]).to("CPU")
@@ -1274,16 +1276,16 @@ class TestCopyFolding(unittest.TestCase):
     check_schedule(x, 3, filter_sink=False)
 
   def test_const_copy_multi(self):
-    x = Tensor.ones(1, device="CPU").to_(["CPU", "CPU:1"])
-    check_schedule(x, 0, filter_sink=False)
-    self.assertEqual(x.item(), 1)
+    x = Tensor.ones(1, device="CPU").to_(["CPU", "CPU:1"]) * 2
+    run_schedule(check_schedule(x, 0, filter_sink=False))
+    self.assertEqual(x.item(), 2.0)
 
   def test_late_const_copy_folding(self):
     a = Tensor.arange(3).realize()
     zeros = Tensor.zeros(3).realize()
-    b = (a*zeros).to("CPU")
+    b = (a*zeros).to("CPU") + 1
     run_schedule(check_schedule(b, 0, filter_sink=False))
-    self.assertListEqual(b.tolist(), [0, 0, 0])
+    self.assertListEqual(b.tolist(), [1, 1, 1])
     self.assertEqual(b.device, "CPU")
 
   def test_alu_after_copy(self):

test/null/test_opts.py

@@ -1,21 +0,0 @@
-import unittest
-from tinygrad import Tensor, Device
-from tinygrad.helpers import CPU_LLVM, CPU_LVP
-from tinygrad.codegen.opt import Opt, OptOps
-from tinygrad.engine.realize import get_program
-
-class TestOpts(unittest.TestCase):
-  def test_opt_upcast(self):
-    opts = (Opt(OptOps.UPCAST, 0, 4),)
-    a = Tensor.empty(16)
-    b = Tensor.empty(16)
-    out = (a+b).contiguous(arg=opts)
-    s = out.schedule()
-    self.assertEqual(s[-1].ast.arg.opts_to_apply, opts)
-    if Device.DEFAULT in {"CPU", "CL", "METAL"} and not CPU_LLVM and not CPU_LVP:
-      prg = get_program(s[-1].ast, renderer=Device[Device.DEFAULT].renderer)
-      self.assertIn('float4', prg.src)
-
-if __name__ == '__main__':
-  unittest.main()
-

test/null/test_schedule.py

@@ -74,7 +74,7 @@ class TestBufferUOp(unittest.TestCase):
     self.assertIsNotNone(a.uop.buffer)
 
   def test_const_does_not_realize(self):
-    a = Tensor(1)+Tensor(2)
+    a = Tensor(1)
     run_schedule(check_schedule(a, 0))
     self.assertIsNone(a.uop.base.realized)
 
@@ -642,6 +642,7 @@ class TestSchedule(unittest.TestCase):
     t = Tensor([1.0, 2.0, 3.0]) ** 8
     self.assertEqual(self._alu_from_tensor(t), [Ops.MUL, Ops.MUL, Ops.MUL])
 
+  @unittest.skip("const folding is removed")
   def test_pow_const_tensor_to_zero(self):
     x = Tensor([1,2,3,4])
     out = x ** Tensor(0.0)
@@ -1003,6 +1004,7 @@ class TestUOpBecome(unittest.TestCase):
 
   # sometimes we prefer to perform an op before movement ops, in this case we should stack the mops on top of the new buffer
 
+  @unittest.skip("no longer supported")
   def test_reorder_expand(self):
     a = Tensor.empty(4, 1)
     b = a.expand(4, 4).reciprocal()
@@ -1038,6 +1040,7 @@ class TestUOpBecome(unittest.TestCase):
     late_add = noop+2
     late_add.realize()
 
+  @unittest.skip("const folding is removed")
   def test_become_const_in_base(self):
     a = Tensor.empty(4)
     b = a*0
@@ -1045,6 +1048,7 @@ class TestUOpBecome(unittest.TestCase):
     check_schedule(b, 0)
     assert UPat(Ops.CONST, arg=0).match(b.uop.base, {}) # scheduling replaces the tensor uop with a VIEW(BUFFER)
 
+  @unittest.skip("const folding is removed")
   def test_become_const_from_const(self):
     const_add = Tensor(1)+Tensor(2)
     assert UPat(Ops.ADD).match(const_add.uop, {})

test/null/test_tensor_uop_representation.py

@@ -8,6 +8,7 @@ def is_pattern_uop(u:UOp, pat:UPat): assert pat.match(u, {}), f"{u}\nis not\n{pa
 def is_pattern(ten:Tensor, pat:UPat): is_pattern_uop(ten.uop, pat)
 
 class TestTensorMutates(unittest.TestCase):
+  @unittest.skip("this doesn't mutate anymore")
   def test_mutate_add(self):
     a = Tensor([1,2,3])
     b = Tensor([4,5,6])

test/null/test_uops_stats.py

@@ -45,6 +45,7 @@ class TestMemoryCount(unittest.TestCase):
     _, mem = get_stats(a+b)
     self.assertEqual(mem, 1024*1024*2 + 1024)  # 1 full read + 1 lil read + 1 write
 
+  @unittest.skip("no longer supported")
   def test_both_expanded(self):
     # TODO: this probably should be a full write
     a = Tensor.empty(1024, 1, dtype=dtypes.uint8).expand(1024, 1024)

tinygrad/uop/ops.py

@@ -506,6 +506,8 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
 
   @functools.cached_property
   def axis(self) -> int|None:
+     # COPY removes axis. TODO: add more tests for this, and consider MSELECT/MSTACK
+    if self.op is Ops.COPY: return None
     if self.op is Ops.MULTI: return self.arg
     # NOTE: they all have to share an axis, we always choose [-1]
     if self.op in GroupOp.ALU: return axes[-1] if (axes := dedup([x.axis for x in self.src if x.axis is not None])) else None