diff --git a/test/test_winograd.py b/test/test_winograd.py index 975fe88b69..e4ef3caa41 100644 --- a/test/test_winograd.py +++ b/test/test_winograd.py @@ -1,10 +1,26 @@ import unittest -from tinygrad import Tensor, GlobalCounters, dtypes +import numpy as np +from tinygrad import Tensor, GlobalCounters, dtypes, Context, nn from tinygrad.uop.ops import Ops from tinygrad.helpers import Timing, CI, Profiling, WINO, DEBUG, getenv from tinygrad.codegen.kernel import Kernel from tinygrad.codegen.heuristic import hand_coded_optimizations +class TestWinogradClose(unittest.TestCase): + def test_close(self): + inp = Tensor.rand(1, 16, 16, 16) + conv = nn.Conv2d(16, 16, 3) + conv(inp).realize() # warmup + GlobalCounters.reset() + print("non winograd") + with Context(WINO=0): + cmp = conv(inp).realize() # warmup + GlobalCounters.reset() + print("winograd") + with Context(WINO=1): + test = conv(inp).realize() + np.testing.assert_allclose(cmp.numpy(), test.numpy(), atol=1e-5) + class TestWinograd(unittest.TestCase): def setUp(self): self.old = WINO.value