more test cleanups (#11544)

* more test cleanups

* revert that

test/external/external_benchmark_schedule.py

@@ -1,12 +1,8 @@
-from typing import List
 from extra.models.resnet import ResNet50
-from tinygrad import Tensor, nn
-from tinygrad.helpers import Profiling, Timing, getenv, BEAM, NOOPT, DEBUG, Context, ansilen
+from tinygrad import Tensor, nn, Device
+from tinygrad.helpers import Profiling, Timing, getenv
 from tinygrad.uop.ops import Ops
-from tinygrad.opt.kernel import Kernel
-from tinygrad.opt.heuristic import hand_coded_optimizations
 from tinygrad.codegen import get_rewrites_for_renderer, apply_rewrites, rewrites_for_linearizer
-from tinygrad.opt.search import beam_search, bufs_from_lin
 from tinygrad.uop.spec import type_verify
 
 if __name__ == "__main__":
@@ -31,26 +27,13 @@ if __name__ == "__main__":
       if not SCHEDULE_ONLY:
         asts = list({x.ast.key:x.ast for x in sched if x.ast.op is Ops.SINK}.values())
         if (restrict_kernel := getenv("RESTRICT_KERNEL", -1)) != -1: asts = asts[restrict_kernel:restrict_kernel+1]
-        kernels: List[Kernel] = []
-        with Timing(f"***** model opts({len(asts):2d}) in  "):
-          with Profiling(PROFILE >= 3):
-            for ast in asts:
-              k = Kernel(ast)
-              if BEAM:
-                with Context(DEBUG=max(2, DEBUG.value)): k = beam_search(k, bufs_from_lin(k), BEAM.value)
-              elif NOOPT: pass
-              else: k.apply_opts(hand_coded_optimizations(k))
-              kernels.append(k)
-
-        with Timing("***** model prep in      "):
-          kernels = [(k, k.get_optimized_ast(), get_rewrites_for_renderer(k.opts, linearizer=False)) for k in kernels]
 
+        rewrites = get_rewrites_for_renderer(Device.default.renderer, linearizer=False)
         with Profiling(PROFILE, fn="/tmp/rewrite.prof"):
           with Timing("***** model rewrite in   "):
             rewritten_uops = []
-            for i,(k,u,rewrites) in enumerate(kernels):
-              with Timing(f"rewrite {i:2d} {k.name}{' '*(50-ansilen(k.name))}", enabled=getenv("VERBOSE", 0)):
-                rewritten_uops.append(apply_rewrites(u, rewrites))
+            for u in asts:
+              rewritten_uops.append(apply_rewrites(u, rewrites))
 
         if LINEARIZE:
           with Timing("***** model linearize in "):

test/test_arange.py

@@ -14,10 +14,7 @@ class TestArange(unittest.TestCase):
     tt = Tensor.arange(N)
     sched = tt.schedule()
     self.assertEqual(len(sched), 1)
-    k = Kernel(sched[-1].ast)
-    if opts is not None:
-      for o in opts: k.apply_opt(o)
-    p = get_program(k.get_optimized_ast(), k.opts)
+    p = get_program(sched[-1].ast, opts=opts)
     print(p.name)
     #print(p.src)
     ExecItem(CompiledRunner(p), [tt.uop.buffer]).run()

test/unit/test_conv_shapetracker.py

@@ -1,55 +0,0 @@
-#!/usr/bin/env python
-import unittest
-from tinygrad.uop.ops import Ops
-from tinygrad.tensor import Tensor
-from tinygrad.nn import Conv2d
-from tinygrad.shape.shapetracker import ShapeTracker, View
-from tinygrad.helpers import prod
-from test.unit.test_shapetracker import shapetracker_getitem
-
-class TestConvShapetracker(unittest.TestCase):
-  def test_conv_3x3_one_view(self):
-    conv = Conv2d(16, 32, (3, 3))
-    # first run to init the weights, they are scheduled.
-    conv(Tensor.empty(1, 16, 10, 10)).schedule()
-    # run it again to get the kernels
-    sched = [si for si in conv(Tensor.empty(1, 16, 10, 10)).schedule() if si.ast.op is Ops.SINK]
-    assert len(sched) == 1, f"conv should only have one kernel, getting {len(sched)}"
-    for st in [x.st_arg for x in sched[0].ast.toposort() if x.op is Ops.LOAD]:
-      assert len(st.views) == 1
-
-  def test_conv_2x2_backward_one_view(self):
-    X = Tensor.rand(1, 1, 3, 3, requires_grad=True)
-    conv = Conv2d(1, 1, (2, 2), bias=False)
-    conv(X).mean().backward()
-    si = X.grad.schedule()[-1]
-    print(si)
-    ldb = [x for x in si.ast.toposort() if x.op is Ops.LOAD][0]
-    st: ShapeTracker = ldb.st_arg.simplify()
-    print(si.bufs[1].size)
-    self.assertEqual(si.bufs[1].size, st.real_size())
-    for v in st.views: print(v)
-
-    # same st
-    test_st = ShapeTracker((
-      View(shape=(1, 1, 2, 4, 2, 4), strides=(0, 0, 2, 8, 1, 4), offset=0, mask=((0, 1), (0, 1), (0, 2), (0, 2), (0, 2), (0, 2)), contiguous=False),
-      View(shape=(1, 1, 1, 1, 3, 3, 3, 3), strides=(0, 0, 0, 0, 24, 8, 3, 1), offset=0,
-           mask=((0, 1), (0, 1), (0, 1), (0, 1), (0, 2), (0, 3), (0, 2), (0, 3)), contiguous=False)))
-    #test_st = ShapeTracker((
-    #  View(shape=(2,4), strides=(1,4), offset=0, mask=None, contiguous=False),
-    #)).simplify()
-      #View(shape=(1, 1, 2, 4, 2, 4), strides=(0, 0, 2, 8, 1, 4), offset=0, mask=((0, 1), (0, 1), (0, 2), (0, 2), (0, 2), (0, 2)), contiguous=False),
-      #View(shape=(1, 1, 1, 1, 3, 3, 3, 3), strides=(0, 0, 0, 0, 24, 8, 3, 1), offset=0,
-      #     mask=((0, 1), (0, 1), (0, 1), (0, 1), (0, 2), (0, 3), (0, 2), (0, 3)), contiguous=False))).simplify()
-    print("*** new ***")
-    for v in test_st.views: print(v)
-    for i in range(prod(st.shape)):
-      i1, i2 = shapetracker_getitem(st, i), shapetracker_getitem(test_st, i)
-      print(i, i1, i2, si.bufs[1].size, i1==i2)
-      #self.assertEqual(i1, i2)
-
-    with self.assertRaises(AssertionError):
-      assert len(st.views) <= 2
-
-if __name__ == '__main__':
-  unittest.main()