From b68b7dbc2aa01362764de74cb0833976cb97b6d6 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Fri, 6 Jun 2025 16:11:49 -0700 Subject: [PATCH] test winograd is close to normal conv [pr] (#10557) Co-authored-by: chenyu --- test/test_winograd.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) 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