diff --git a/test/opt/test_hand_coded_opts.py b/test/opt/test_hand_coded_opts.py index f4448850d1..8c9e2e2f01 100644 --- a/test/opt/test_hand_coded_opts.py +++ b/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)) diff --git a/test/unit/test_winograd.py b/test/unit/test_winograd.py index 6cd38ab3b2..070f18ade2 100644 --- a/test/unit/test_winograd.py +++ b/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