clean up wino tests (#12008)

removed the one that tests hcopt and added one for backward kernel counts

test/opt/test_hand_coded_opts.py

@@ -1,6 +1,6 @@
 import unittest
 from tinygrad import Tensor, Device
-from tinygrad.helpers import Context, prod
+from tinygrad.helpers import prod
 from tinygrad.uop.ops import AxisType
 from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
 
@@ -33,35 +33,6 @@ class TestHandCodedOpts(unittest.TestCase):
     # should upcast the two Tensor.stacks
     assert k.upcasted >= 2 and k.full_shape[k.shape_len-k.upcasted:k.shape_len].count(6) == 2
 
-  def test_masked_upcast_wino_full(self):
-    with Context(WINO=1):
-      x,w = Tensor.rand(1,4,8,8, requires_grad=True).realize(), Tensor.rand(4,4,3,3, requires_grad=True).realize()
-      out = Tensor.conv2d(x,w, padding=1)
-      out.mean().backward()
-
-      upcasts = []
-      wino_schedule = out.schedule()
-      # collect upcasts of tile transform kernels
-      for i, si in enumerate(wino_schedule):
-        k = Kernel(push_views(si.ast))
-        k.apply_opts(hand_coded_optimizations(k))
-        if k.reduceop is not None: continue  # not a tile transform kernel (there is a gemm reduce kernel)
-        if len(k.bufs) < 22: continue  # not a tile transform kernel (there's a permute kernel at the end)
-        upcasts.append(tuple(k.full_shape[k.shape_len - k.upcasted:k.shape_len]))
-      assert len(upcasts) == 3  # 3 transformation matrices
-      assert len(wino_schedule) <= 4  # 4 kernels
-      # this test case's inputs are too small, so one of the 4-stacks became a local, which is fine i guess
-      assert upcasts.count((6, 6)) == 2 #and upcasts.count((4, 4)) == 1
-
-      backward_schedule = Tensor.schedule(x.grad, w.grad)
-      for si in backward_schedule:
-        k = Kernel(push_views(si.ast))
-        k.apply_opts(hand_coded_optimizations(k))
-        if len(k.bufs) < 20: continue  # not a tile transform kernel
-        # heuristic number to make sure that at least some upcasts but not too many upcasts are being done
-        assert 6 <= prod(k.full_shape[k.shape_len - k.upcasted:k.shape_len]) <= 216
-      assert len(backward_schedule) <= 13  # just the current number, but it could be better
-
   def test_masked_upcast_many(self):
     layer_1 = Tensor.cat(Tensor.rand(3, 4), Tensor.rand(4, 4))
     layer_2 = Tensor.cat(layer_1.unsqueeze(0), Tensor.rand(6, 7, 4))

test/unit/test_winograd.py

@@ -28,15 +28,19 @@ class TestWinograd(unittest.TestCase):
   def test_profile(self):
     x,w = Tensor.rand(1,4,9,9).realize(), Tensor.rand(4,4,3,3).realize()
     with Profiling(enabled=not CI, sort='time'):
-      out = Tensor.conv2d(x,w).realize()
-    out.numpy()
+      Tensor.conv2d(x,w).realize()
 
-  def test_four_kernels(self):
+  def test_forward_kernels(self):
     x,w = Tensor.rand(1,4,9,9).realize(), Tensor.rand(4,4,3,3).realize()
-    GlobalCounters.reset()
-    out = Tensor.conv2d(x,w).realize()
-    assert GlobalCounters.kernel_count == 4
-    out.numpy()
+    out = Tensor.conv2d(x,w)
+    self.assertEqual(len(out.schedule()), 4)
+
+  def test_backward_kernels(self):
+    x,w = Tensor.empty(1,4,9,9,requires_grad=True).realize(), Tensor.empty(4,4,3,3,requires_grad=True).realize()
+    out = Tensor.conv2d(x,w, padding=1)
+    out.mean().backward()
+    backward_schedule = Tensor.schedule(x.grad, w.grad)
+    self.assertEqual(len(backward_schedule), 9)
 
   def test_counters(self):
     IC, OC, X, Y = 4,4,9,9